





'''simple docstring'''






import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)
def 							_lowerCamelCase(  UpperCamelCase__    : str    ,       UpperCamelCase__    : List[Any]=False		)		->  List[Any]:
    A       :       Tuple							=						[]
    for i in range(config.num_hidden_layers		):
        # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
        rename_keys.append((F'''module.blocks.{i}.norm1.weight''', F'''vit.encoder.layer.{i}.layernorm_before.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.norm1.bias''', F'''vit.encoder.layer.{i}.layernorm_before.bias''')		)
        rename_keys.append(
            (F'''module.blocks.{i}.attn.proj.weight''', F'''vit.encoder.layer.{i}.attention.output.dense.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.attn.proj.bias''', F'''vit.encoder.layer.{i}.attention.output.dense.bias''')		)
        rename_keys.append((F'''module.blocks.{i}.norm2.weight''', F'''vit.encoder.layer.{i}.layernorm_after.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.norm2.bias''', F'''vit.encoder.layer.{i}.layernorm_after.bias''')		)
        rename_keys.append((F'''module.blocks.{i}.mlp.fc1.weight''', F'''vit.encoder.layer.{i}.intermediate.dense.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.mlp.fc1.bias''', F'''vit.encoder.layer.{i}.intermediate.dense.bias''')		)
        rename_keys.append((F'''module.blocks.{i}.mlp.fc2.weight''', F'''vit.encoder.layer.{i}.output.dense.weight''')		)
        rename_keys.append((F'''module.blocks.{i}.mlp.fc2.bias''', F'''vit.encoder.layer.{i}.output.dense.bias''')		)

    # projection layer + position embeddings
    rename_keys.extend(
        [
            ('''module.cls_token''', '''vit.embeddings.cls_token'''),
            ('''module.patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
            ('''module.patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
            ('''module.pos_embed''', '''vit.embeddings.position_embeddings'''),
        ]		)

    if base_model:
        # layernorm + pooler
        rename_keys.extend(
            [
                ('''module.norm.weight''', '''layernorm.weight'''),
                ('''module.norm.bias''', '''layernorm.bias'''),
            ]		)

        # if just the base model, we should remove "vit" from all keys that start with "vit"
        A       :       Any							=						[(pair[0], pair[1][4:]) if pair[1].startswith('''vit'''		) else pair for pair in rename_keys]
    else:
        # layernorm + classification head
        rename_keys.extend(
            [
                ('''norm.weight''', '''vit.layernorm.weight'''),
                ('''norm.bias''', '''vit.layernorm.bias'''),
                ('''head.weight''', '''classifier.weight'''),
                ('''head.bias''', '''classifier.bias'''),
            ]		)

    return rename_keys
def 							_lowerCamelCase(  UpperCamelCase__    : Tuple    ,       UpperCamelCase__    : Optional[Any]    ,       UpperCamelCase__    : str=False		)		->  str:
    for i in range(config.num_hidden_layers		):
        if base_model:
            A       :       Dict							=						''''''
        else:
            A       :       Optional[int]							=						'''vit.'''
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        A       :       Optional[Any]							=						state_dict.pop(F'''module.blocks.{i}.attn.qkv.weight'''		)
        A       :       str							=						state_dict.pop(F'''module.blocks.{i}.attn.qkv.bias'''		)
        # next, add query, keys and values (in that order) to the state dict
        A       :       int							=						in_proj_weight[
            : config.hidden_size, :
        ]
        A       :       List[str]							=						in_proj_bias[: config.hidden_size]
        A       :       str							=						in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        A       :       Tuple							=						in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        A       :       Optional[Any]							=						in_proj_weight[
            -config.hidden_size :, :
        ]
        A       :       Tuple							=						in_proj_bias[-config.hidden_size :]
def 							_lowerCamelCase(  UpperCamelCase__    : str		)		->  Optional[int]:
    A       :       Optional[Any]							=						['''head.weight''', '''head.bias''']
    for k in ignore_keys:
        state_dict.pop(UpperCamelCase__    ,       UpperCamelCase__		)
def 							_lowerCamelCase(  UpperCamelCase__    : Optional[int]		)		->  str:
    # projection head is used in the self-supervised pre-training in MSN,
    # for downstream task it's not needed.
    A       :       List[Any]							=						[
        '''module.fc.fc1.weight''',
        '''module.fc.fc1.bias''',
        '''module.fc.bn1.weight''',
        '''module.fc.bn1.bias''',
        '''module.fc.bn1.running_mean''',
        '''module.fc.bn1.running_var''',
        '''module.fc.bn1.num_batches_tracked''',
        '''module.fc.fc2.weight''',
        '''module.fc.fc2.bias''',
        '''module.fc.bn2.weight''',
        '''module.fc.bn2.bias''',
        '''module.fc.bn2.running_mean''',
        '''module.fc.bn2.running_var''',
        '''module.fc.bn2.num_batches_tracked''',
        '''module.fc.fc3.weight''',
        '''module.fc.fc3.bias''',
    ]
    for k in ignore_keys:
        state_dict.pop(UpperCamelCase__    ,       UpperCamelCase__		)
def 							_lowerCamelCase(  UpperCamelCase__    : Union[str, Any]    ,       UpperCamelCase__    : Tuple    ,       UpperCamelCase__    : str		)		->  Optional[Any]:
    A       :       List[Any]							=						dct.pop(UpperCamelCase__		)
    A       :       Optional[int]							=						val
def 							_lowerCamelCase(  UpperCamelCase__    : Optional[Any]    ,       UpperCamelCase__    : Optional[Any]		)		->  List[str]:
    A       :       Optional[Any]							=						ViTMSNConfig()
    A       :       Optional[Any]							=						1_000

    A       :       Tuple							=						'''datasets/huggingface/label-files'''
    A       :       List[str]							=						'''imagenet-1k-id2label.json'''
    A       :       Any							=						json.load(open(hf_hub_download(UpperCamelCase__    ,       UpperCamelCase__		)    ,       '''r'''		)		)
    A       :       Any							=						{int(UpperCamelCase__		): v for k, v in idalabel.items()}
    A       :       str							=						idalabel
    A       :       Any							=						{v: k for k, v in idalabel.items()}

    if "s16" in checkpoint_url:
        A       :       Union[str, Any]							=						384
        A       :       Any							=						1_536
        A       :       Any							=						6
    elif "l16" in checkpoint_url:
        A       :       Dict							=						1_024
        A       :       Tuple							=						4_096
        A       :       int							=						24
        A       :       Dict							=						16
        A       :       int							=						0.1
    elif "b4" in checkpoint_url:
        A       :       str							=						4
    elif "l7" in checkpoint_url:
        A       :       Optional[int]							=						7
        A       :       Optional[int]							=						1_024
        A       :       int							=						4_096
        A       :       str							=						24
        A       :       Any							=						16
        A       :       Union[str, Any]							=						0.1

    A       :       Dict							=						ViTMSNModel(UpperCamelCase__		)

    A       :       Union[str, Any]							=						torch.hub.load_state_dict_from_url(UpperCamelCase__    ,       map_location='''cpu'''		)['''target_encoder''']

    A       :       int							=						ViTImageProcessor(size=config.image_size		)

    remove_projection_head(UpperCamelCase__		)
    A       :       Tuple							=						create_rename_keys(UpperCamelCase__    ,       base_model=UpperCamelCase__		)

    for src, dest in rename_keys:
        rename_key(UpperCamelCase__    ,       UpperCamelCase__    ,       UpperCamelCase__		)
    read_in_q_k_v(UpperCamelCase__    ,       UpperCamelCase__    ,       base_model=UpperCamelCase__		)

    model.load_state_dict(UpperCamelCase__		)
    model.eval()

    A       :       Union[str, Any]							=						'''http://images.cocodataset.org/val2017/000000039769.jpg'''

    A       :       Optional[Any]							=						Image.open(requests.get(UpperCamelCase__    ,       stream=UpperCamelCase__		).raw		)
    A       :       Dict							=						ViTImageProcessor(
        size=config.image_size    ,       image_mean=UpperCamelCase__    ,       image_std=UpperCamelCase__		)
    A       :       List[str]							=						image_processor(images=UpperCamelCase__    ,       return_tensors='''pt'''		)

    # forward pass
    torch.manual_seed(2		)
    A       :       Optional[Any]							=						model(**UpperCamelCase__		)
    A       :       int							=						outputs.last_hidden_state

    # The following Colab Notebook was used to generate these outputs:
    # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
    if "s16" in checkpoint_url:
        A       :       str							=						torch.tensor([[-1.0_9_1_5, -1.4_8_7_6, -1.1_8_0_9]]		)
    elif "b16" in checkpoint_url:
        A       :       int							=						torch.tensor([[1_4.2_8_8_9, -1_8.9_0_4_5, 1_1.7_2_8_1]]		)
    elif "l16" in checkpoint_url:
        A       :       Optional[Any]							=						torch.tensor([[4_1.5_0_2_8, -2_2.8_6_8_1, 4_5.6_4_7_5]]		)
    elif "b4" in checkpoint_url:
        A       :       Optional[Any]							=						torch.tensor([[-4.3_8_6_8, 5.2_9_3_2, -0.4_1_3_7]]		)
    else:
        A       :       str							=						torch.tensor([[-0.1_7_9_2, -0.6_4_6_5, 2.4_2_6_3]]		)

    # verify logits
    assert torch.allclose(last_hidden_state[:, 0, :3]    ,       UpperCamelCase__    ,       atol=1e-4		)

    print(F'''Saving model to {pytorch_dump_folder_path}'''		)
    model.save_pretrained(UpperCamelCase__		)

    print(F'''Saving image processor to {pytorch_dump_folder_path}'''		)
    image_processor.save_pretrained(UpperCamelCase__		)


if __name__ == "__main__":
     snake_case_     =	argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         """--checkpoint_url""",
         default="""https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar""",
         type=str,
         help="""URL of the checkpoint you'd like to convert.""",
     )
     parser.add_argument(
         """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
     )

     snake_case_     =	parser.parse_args()
     convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)







'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_torch_available,
)


snake_case_     =	{
    """configuration_speecht5""": [
        """SPEECHT5_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """SPEECHT5_PRETRAINED_HIFIGAN_CONFIG_ARCHIVE_MAP""",
        """SpeechT5Config""",
        """SpeechT5HifiGanConfig""",
    ],
    """feature_extraction_speecht5""": ["""SpeechT5FeatureExtractor"""],
    """processing_speecht5""": ["""SpeechT5Processor"""],
}

try:
     if not is_sentencepiece_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case_     =	["""SpeechT5Tokenizer"""]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case_     =	[
         """SPEECHT5_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """SpeechT5ForSpeechToText""",
         """SpeechT5ForSpeechToSpeech""",
         """SpeechT5ForTextToSpeech""",
         """SpeechT5Model""",
         """SpeechT5PreTrainedModel""",
         """SpeechT5HifiGan""",
     ]

if TYPE_CHECKING:
     from .configuration_speechta import (
         SPEECHT5_PRETRAINED_CONFIG_ARCHIVE_MAP,
         SPEECHT5_PRETRAINED_HIFIGAN_CONFIG_ARCHIVE_MAP,
         SpeechTaConfig,
         SpeechTaHifiGanConfig,
     )
     from .feature_extraction_speechta import SpeechTaFeatureExtractor
     from .processing_speechta import SpeechTaProcessor

     try:
          if not is_sentencepiece_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_speechta import SpeechTaTokenizer

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_speechta import (
              SPEECHT5_PRETRAINED_MODEL_ARCHIVE_LIST,
              SpeechTaForSpeechToSpeech,
              SpeechTaForSpeechToText,
              SpeechTaForTextToSpeech,
              SpeechTaHifiGan,
              SpeechTaModel,
              SpeechTaPreTrainedModel,
          )

else:
     import sys

     snake_case_     =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
