




'''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     SCREAMING_SNAKE_CASE							(   a_						:				List[Any]					,       a_						:				List[str]=False      ):
  __a								=       []
  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								=       [(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     SCREAMING_SNAKE_CASE							(   a_						:				Any					,       a_						:				List[Any]					,       a_						:				Dict=False      ):
  for i in range(config.num_hidden_layers      ):
    if base_model:
      __a								=       ''
    else:
      __a								=       'vit.'
    # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
    __a								=       state_dict.pop(f"module.blocks.{i}.attn.qkv.weight"      )
    __a								=       state_dict.pop(f"module.blocks.{i}.attn.qkv.bias"      )
    # next, add query, keys and values (in that order) to the state dict
    __a								=       in_proj_weight[
        : config.hidden_size, :
    ]
    __a								=       in_proj_bias[: config.hidden_size]
    __a								=       in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    __a								=       in_proj_bias[
        config.hidden_size : config.hidden_size * 2
    ]
    __a								=       in_proj_weight[
        -config.hidden_size :, :
    ]
    __a								=       in_proj_bias[-config.hidden_size :]


def     SCREAMING_SNAKE_CASE							(   a_						:				Optional[int]      ):
  __a								=       ['head.weight', 'head.bias']
  for k in ignore_keys:
    state_dict.pop(a_					,       a_      )


def     SCREAMING_SNAKE_CASE							(   a_						:				Tuple      ):
  # projection head is used in the self-supervised pre-training in MSN,
  # for downstream task it's not needed.
  __a								=       [
      '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(a_					,       a_      )


def     SCREAMING_SNAKE_CASE							(   a_						:				Union[str, Any]					,       a_						:				str					,       a_						:				List[Any]      ):
  __a								=       dct.pop(a_      )
  __a								=       val


def     SCREAMING_SNAKE_CASE							(   a_						:				Union[str, Any]					,       a_						:				List[str]      ):
  __a								=       ViTMSNConfig()
  __a								=       1000

  __a								=       'datasets/huggingface/label-files'
  __a								=       'imagenet-1k-id2label.json'
  __a								=       json.load(open(hf_hub_download(a_					,       a_      )					,       'r'      )      )
  __a								=       {int(a_      ): v for k, v in idalabel.items()}
  __a								=       idalabel
  __a								=       {v: k for k, v in idalabel.items()}

  if "s16" in checkpoint_url:
    __a								=       384
    __a								=       1536
    __a								=       6
  elif "l16" in checkpoint_url:
    __a								=       1024
    __a								=       4096
    __a								=       24
    __a								=       16
    __a								=       0.1
  elif "b4" in checkpoint_url:
    __a								=       4
  elif "l7" in checkpoint_url:
    __a								=       7
    __a								=       1024
    __a								=       4096
    __a								=       24
    __a								=       16
    __a								=       0.1

  __a								=       ViTMSNModel(a_      )

  __a								=       torch.hub.load_state_dict_from_url(a_					,       map_location='cpu'      )['target_encoder']

  __a								=       ViTImageProcessor(size=config.image_size      )

  remove_projection_head(a_      )
  __a								=       create_rename_keys(a_					,       base_model=a_      )

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

  model.load_state_dict(a_      )
  model.eval()

  __a								=       'http://images.cocodataset.org/val2017/000000039769.jpg'

  __a								=       Image.open(requests.get(a_					,       stream=a_      ).raw      )
  __a								=       ViTImageProcessor(
      size=config.image_size					,       image_mean=a_					,       image_std=a_      )
  __a								=       image_processor(images=a_					,       return_tensors='pt'      )

  # forward pass
  torch.manual_seed(2      )
  __a								=       model(**a_      )
  __a								=       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								=       torch.tensor([[-1.09_15, -1.48_76, -1.18_09]]      )
  elif "b16" in checkpoint_url:
    __a								=       torch.tensor([[14.28_89, -18.90_45, 11.72_81]]      )
  elif "l16" in checkpoint_url:
    __a								=       torch.tensor([[41.50_28, -22.86_81, 45.64_75]]      )
  elif "b4" in checkpoint_url:
    __a								=       torch.tensor([[-4.38_68, 5.29_32, -0.41_37]]      )
  else:
    __a								=       torch.tensor([[-0.17_92, -0.64_65, 2.42_63]]      )

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

  print(f"Saving model to {pytorch_dump_folder_path}"      )
  model.save_pretrained(a_      )

  print(f"Saving image processor to {pytorch_dump_folder_path}"      )
  image_processor.save_pretrained(a_      )


if __name__ == "__main__":
  UpperCAmelCase_					=   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."
  )

  UpperCAmelCase_					=   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_torch_available


UpperCAmelCase_					=   {
    "configuration_clap": [
        "CLAP_PRETRAINED_MODEL_ARCHIVE_LIST",
        "ClapAudioConfig",
        "ClapConfig",
        "ClapTextConfig",
    ],
    "processing_clap": ["ClapProcessor"],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_					=   [
      "CLAP_PRETRAINED_MODEL_ARCHIVE_LIST",
      "ClapModel",
      "ClapPreTrainedModel",
      "ClapTextModel",
      "ClapTextModelWithProjection",
      "ClapAudioModel",
      "ClapAudioModelWithProjection",
  ]
  UpperCAmelCase_					=   ["ClapFeatureExtractor"]

if TYPE_CHECKING:
  from .configuration_clap import (
      CLAP_PRETRAINED_MODEL_ARCHIVE_LIST,
      ClapAudioConfig,
      ClapConfig,
      ClapTextConfig,
  )
  from .processing_clap import ClapProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .feature_extraction_clap import ClapFeatureExtractor
    from .modeling_clap import (
        CLAP_PRETRAINED_MODEL_ARCHIVE_LIST,
        ClapAudioModel,
        ClapAudioModelWithProjection,
        ClapModel,
        ClapPreTrainedModel,
        ClapTextModel,
        ClapTextModelWithProjection,
    )


else:
  import sys

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



