




'''simple docstring'''





import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor


def      UpperCamelCase_(	snake_case							:       Any     ):

							'''simple docstring'''






							if "cls_token" in name:
														snake_case_						=      name.replace("cls_token"							,				"vit.embeddings.cls_token"     )
							if "mask_token" in name:
														snake_case_						=      name.replace("mask_token"							,				"decoder.mask_token"     )
							if "decoder_pos_embed" in name:
														snake_case_						=      name.replace("decoder_pos_embed"							,				"decoder.decoder_pos_embed"     )
							if "pos_embed" in name and "decoder" not in name:
														snake_case_						=      name.replace("pos_embed"							,				"vit.embeddings.position_embeddings"     )
							if "patch_embed.proj" in name:
														snake_case_						=      name.replace("patch_embed.proj"							,				"vit.embeddings.patch_embeddings.projection"     )
							if "patch_embed.norm" in name:
														snake_case_						=      name.replace("patch_embed.norm"							,				"vit.embeddings.norm"     )
							if "decoder_blocks" in name:
														snake_case_						=      name.replace("decoder_blocks"							,				"decoder.decoder_layers"     )
							if "blocks" in name:
														snake_case_						=      name.replace("blocks"							,				"vit.encoder.layer"     )
							if "attn.proj" in name:
														snake_case_						=      name.replace("attn.proj"							,				"attention.output.dense"     )
							if "attn" in name:
														snake_case_						=      name.replace("attn"							,				"attention.self"     )
							if "norm1" in name:
														snake_case_						=      name.replace("norm1"							,				"layernorm_before"     )
							if "norm2" in name:
														snake_case_						=      name.replace("norm2"							,				"layernorm_after"     )
							if "mlp.fc1" in name:
														snake_case_						=      name.replace("mlp.fc1"							,				"intermediate.dense"     )
							if "mlp.fc2" in name:
														snake_case_						=      name.replace("mlp.fc2"							,				"output.dense"     )
							if "decoder_embed" in name:
														snake_case_						=      name.replace("decoder_embed"							,				"decoder.decoder_embed"     )
							if "decoder_norm" in name:
														snake_case_						=      name.replace("decoder_norm"							,				"decoder.decoder_norm"     )
							if "decoder_pred" in name:
														snake_case_						=      name.replace("decoder_pred"							,				"decoder.decoder_pred"     )
							if "norm.weight" in name and "decoder" not in name:
														snake_case_						=      name.replace("norm.weight"							,				"vit.layernorm.weight"     )
							if "norm.bias" in name and "decoder" not in name:
														snake_case_						=      name.replace("norm.bias"							,				"vit.layernorm.bias"     )

							return name


def      UpperCamelCase_(	snake_case							:       Optional[Any]							,				snake_case							:       Any     ):

							'''simple docstring'''






							for key in orig_state_dict.copy().keys():
														snake_case_						=      orig_state_dict.pop(SCREAMING_SNAKE_CASE__     )

														if "qkv" in key:
																					snake_case_						=      key.split("."     )
																					snake_case_						=      int(key_split[1]     )
																					if "decoder_blocks" in key:
																												snake_case_						=      config.decoder_hidden_size
																												snake_case_						=      """decoder.decoder_layers."""
																												if "weight" in key:
																																			snake_case_						=      val[:dim, :]
																																			snake_case_						=      val[dim : dim * 2, :]
																																			snake_case_						=      val[-dim:, :]
																												elif "bias" in key:
																																			snake_case_						=      val[:dim]
																																			snake_case_						=      val[dim : dim * 2]
																																			snake_case_						=      val[-dim:]
																					else:
																												snake_case_						=      config.hidden_size
																												snake_case_						=      """vit.encoder.layer."""
																												if "weight" in key:
																																			snake_case_						=      val[:dim, :]
																																			snake_case_						=      val[dim : dim * 2, :]
																																			snake_case_						=      val[-dim:, :]
																												elif "bias" in key:
																																			snake_case_						=      val[:dim]
																																			snake_case_						=      val[dim : dim * 2]
																																			snake_case_						=      val[-dim:]

														else:
																					snake_case_						=      val

							return orig_state_dict


def      UpperCamelCase_(	snake_case							:       Union[str, Any]							,				snake_case							:       str     ):

							'''simple docstring'''






							snake_case_						=      ViTMAEConfig()
							if "large" in checkpoint_url:
														snake_case_						=      1_0_2_4
														snake_case_						=      4_0_9_6
														snake_case_						=      2_4
														snake_case_						=      1_6
							elif "huge" in checkpoint_url:
														snake_case_						=      1_4
														snake_case_						=      1_2_8_0
														snake_case_						=      5_1_2_0
														snake_case_						=      3_2
														snake_case_						=      1_6

							snake_case_						=      ViTMAEForPreTraining(SCREAMING_SNAKE_CASE__     )

							snake_case_						=      torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE__							,				map_location="cpu"     )["""model"""]

							snake_case_						=      ViTMAEImageProcessor(size=config.image_size     )

							snake_case_						=      convert_state_dict(SCREAMING_SNAKE_CASE__							,				SCREAMING_SNAKE_CASE__     )

							model.load_state_dict(SCREAMING_SNAKE_CASE__     )
							model.eval()

							snake_case_						=      """https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg"""

							snake_case_						=      Image.open(requests.get(SCREAMING_SNAKE_CASE__							,				stream=SCREAMING_SNAKE_CASE__     ).raw     )
							snake_case_						=      ViTMAEImageProcessor(size=config.image_size     )
							snake_case_						=      image_processor(images=SCREAMING_SNAKE_CASE__							,				return_tensors="pt"     )

							# forward pass
							torch.manual_seed(2     )
							snake_case_						=      model(**SCREAMING_SNAKE_CASE__     )
							snake_case_						=      outputs.logits

							if "large" in checkpoint_url:
														snake_case_						=      torch.tensor(
														    [[-0.7309, -0.7128, -1.0169], [-1.0161, -0.9058, -1.1878], [-1.0478, -0.9411, -1.1911]]     )
							elif "huge" in checkpoint_url:
														snake_case_						=      torch.tensor(
														    [[-1.1599, -0.9199, -1.2221], [-1.1952, -0.9269, -1.2307], [-1.2143, -0.9337, -1.2262]]     )
							else:
														snake_case_						=      torch.tensor(
														    [[-0.9192, -0.8481, -1.1259], [-1.1349, -1.0034, -1.2599], [-1.1757, -1.0429, -1.2726]]     )

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

							print(f'Saving model to {pytorch_dump_folder_path}'     )
							model.save_pretrained(SCREAMING_SNAKE_CASE__     )

							print(f'Saving image processor to {pytorch_dump_folder_path}'     )
							image_processor.save_pretrained(SCREAMING_SNAKE_CASE__     )


if __name__ == "__main__":
		_SCREAMING_SNAKE_CASE							:   Union[str, Any]      		= argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--checkpoint_url",
		    default="https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth",
		    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."
		)

		_SCREAMING_SNAKE_CASE							:   List[Any]      		= parser.parse_args()
		convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






from collections import UserDict
from typing import Union

import numpy as np
import requests

from ..utils import (
    add_end_docstrings,
    logging,
)
from .audio_classification import ffmpeg_read
from .base import PIPELINE_INIT_ARGS, Pipeline


lowercase  :					Optional[Any]						    =							logging.get_logger(__name__)

@add_end_docstrings(lowerCAmelCase     )
class __snake_case		(      lowerCAmelCase     ):
    def __init__(  self		,**snake_case      ):



           '''simple docstring'''







           super().__init__(**snake_case      )

           if self.framework != "pt":
                  raise ValueError(f"The {self.__class__} is only available in PyTorch."      )
           # No specific FOR_XXX available yet
    def __call__(  self		,snake_case		,**snake_case      ):



           '''simple docstring'''







           return super().__call__(snake_case		,**snake_case      )
    def       _SCREAMING_SNAKE_CASE       (  self		,**snake_case      ):



           '''simple docstring'''







           lowercase							:     Union[str, Any]       	=      {}
           if "candidate_labels" in kwargs:
                  lowercase							:     List[str]       	=      kwargs["""candidate_labels"""]
           if "hypothesis_template" in kwargs:
                  lowercase							:     Dict       	=      kwargs["""hypothesis_template"""]

           return preprocess_params, {}, {}
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case		,snake_case=None		,snake_case="This is a sound of {}."      ):



           '''simple docstring'''







           if isinstance(snake_case		,snake_case      ):
                  if audio.startswith("""http://"""      ) or audio.startswith("""https://"""      ):
                         # We need to actually check for a real protocol, otherwise it's impossible to use a local file
                         # like http_huggingface_co.png
                         lowercase							:     Optional[Any]       	=      requests.get(snake_case      ).content
                  else:
                         with open(snake_case		,"""rb"""      ) as f:
                                lowercase							:     Union[str, Any]       	=      f.read()

           if isinstance(snake_case		,snake_case      ):
                  lowercase							:     int       	=      ffmpeg_read(snake_case		,self.feature_extractor.sampling_rate      )

           if not isinstance(snake_case		,np.ndarray      ):
                  raise ValueError("""We expect a numpy ndarray as input"""      )
           if len(audio.shape      ) != 1:
                  raise ValueError("""We expect a single channel audio input for ZeroShotAudioClassificationPipeline"""      )

           lowercase							:     Dict       	=      self.feature_extractor(
               [audio]		,sampling_rate=self.feature_extractor.sampling_rate		,return_tensors="""pt"""      )
           lowercase							:     Tuple       	=      candidate_labels
           lowercase							:     Tuple       	=      [hypothesis_template.format(snake_case      ) for x in candidate_labels]
           lowercase							:     Optional[Any]       	=      self.tokenizer(snake_case		,return_tensors=self.framework		,padding=snake_case      )
           lowercase							:     Optional[Any]       	=      [text_inputs]
           return inputs
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case      ):



           '''simple docstring'''







           lowercase							:     List[str]       	=      model_inputs.pop("""candidate_labels"""      )
           lowercase							:     Dict       	=      model_inputs.pop("""text_inputs"""      )
           if isinstance(text_inputs[0]		,snake_case      ):
                  lowercase							:     List[Any]       	=      text_inputs[0]
           else:
                  # Batching case.
                  lowercase							:     Dict       	=      text_inputs[0][0]

           lowercase							:     Optional[Any]       	=      self.model(**snake_case		,**snake_case      )

           lowercase							:     Any       	=      {
               """candidate_labels""": candidate_labels,
               """logits""": outputs.logits_per_audio,
           }
           return model_outputs
    def       _SCREAMING_SNAKE_CASE       (  self		,snake_case      ):



           '''simple docstring'''







           lowercase							:     List[Any]       	=      model_outputs.pop("""candidate_labels"""      )
           lowercase							:     Any       	=      model_outputs["""logits"""][0]

           if self.framework == "pt":
                  lowercase							:     Any       	=      logits.softmax(dim=0      )
                  lowercase							:     Tuple       	=      probs.tolist()
           else:
                  raise ValueError("""`tf` framework not supported."""      )

           lowercase							:     Tuple       	=      [
               {"""score""": score, """label""": candidate_label}
               for score, candidate_label in sorted(zip(snake_case		,snake_case      )		,key=lambda snake_case      : -x[0]      )
           ]
           return result
