





"""simple docstring"""


import argparse
import json
from pathlib import Path

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

from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


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




def 					lowercase__    (					lowercase_		,lowercase_=False			)					->		Any:




      """simple docstring"""

      _UpperCamelCase      :	Union[str, Any]						=       []
      for i in range(config.num_hidden_layers			):
            # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
            rename_keys.append((F'''blocks.{i}.norm1.weight''', F'''vit.encoder.layer.{i}.layernorm_before.weight''')			)
            rename_keys.append((F'''blocks.{i}.norm1.bias''', F'''vit.encoder.layer.{i}.layernorm_before.bias''')			)
            rename_keys.append((F'''blocks.{i}.attn.proj.weight''', F'''vit.encoder.layer.{i}.attention.output.dense.weight''')			)
            rename_keys.append((F'''blocks.{i}.attn.proj.bias''', F'''vit.encoder.layer.{i}.attention.output.dense.bias''')			)
            rename_keys.append((F'''blocks.{i}.norm2.weight''', F'''vit.encoder.layer.{i}.layernorm_after.weight''')			)
            rename_keys.append((F'''blocks.{i}.norm2.bias''', F'''vit.encoder.layer.{i}.layernorm_after.bias''')			)
            rename_keys.append((F'''blocks.{i}.mlp.fc1.weight''', F'''vit.encoder.layer.{i}.intermediate.dense.weight''')			)
            rename_keys.append((F'''blocks.{i}.mlp.fc1.bias''', F'''vit.encoder.layer.{i}.intermediate.dense.bias''')			)
            rename_keys.append((F'''blocks.{i}.mlp.fc2.weight''', F'''vit.encoder.layer.{i}.output.dense.weight''')			)
            rename_keys.append((F'''blocks.{i}.mlp.fc2.bias''', F'''vit.encoder.layer.{i}.output.dense.bias''')			)

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

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

            # if just the base model, we should remove "vit" from all keys that start with "vit"
            _UpperCamelCase      :	List[str]						=       [(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 					lowercase__    (					lowercase_		,lowercase_		,lowercase_=False			)					->		Any:




      """simple docstring"""

      for i in range(config.num_hidden_layers			):
            if base_model:
                  _UpperCamelCase      :	Union[str, Any]						=       ""
            else:
                  _UpperCamelCase      :	Union[str, Any]						=       "vit."
            # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
            _UpperCamelCase      :	Any						=       state_dict.pop(F'''blocks.{i}.attn.qkv.weight'''			)
            _UpperCamelCase      :	List[str]						=       state_dict.pop(F'''blocks.{i}.attn.qkv.bias'''			)
            # next, add query, keys and values (in that order) to the state dict
            _UpperCamelCase      :	List[Any]						=       in_proj_weight[
                : config.hidden_size, :
            ]
            _UpperCamelCase      :	Tuple						=       in_proj_bias[: config.hidden_size]
            _UpperCamelCase      :	Any						=       in_proj_weight[
                config.hidden_size : config.hidden_size * 2, :
            ]
            _UpperCamelCase      :	Optional[int]						=       in_proj_bias[
                config.hidden_size : config.hidden_size * 2
            ]
            _UpperCamelCase      :	Dict						=       in_proj_weight[
                -config.hidden_size :, :
            ]
            _UpperCamelCase      :	Any						=       in_proj_bias[-config.hidden_size :]




def 					lowercase__    (					lowercase_			)					->		Any:




      """simple docstring"""

      _UpperCamelCase      :	Optional[int]						=       ["head.weight", "head.bias"]
      for k in ignore_keys:
            state_dict.pop(lowercase_		,lowercase_			)




def 					lowercase__    (					lowercase_		,lowercase_		,lowercase_			)					->		List[str]:




      """simple docstring"""

      _UpperCamelCase      :	Dict						=       dct.pop(lowercase_			)
      _UpperCamelCase      :	List[str]						=       val




def 					lowercase__    (					)					->		List[Any]:




      """simple docstring"""

      _UpperCamelCase      :	Union[str, Any]						=       "http://images.cocodataset.org/val2017/000000039769.jpg"
      _UpperCamelCase      :	Any						=       Image.open(requests.get(lowercase_		,stream=lowercase_			).raw			)
      return im




@torch.no_grad()
def 					lowercase__    (					lowercase_		,lowercase_		,lowercase_=True			)					->		List[str]:




      """simple docstring"""

      _UpperCamelCase      :	Union[str, Any]						=       ViTConfig()
      # patch_size
      if model_name[-1] == "8":
            _UpperCamelCase      :	List[Any]						=       8
      # set labels if required
      if not base_model:
            _UpperCamelCase      :	Tuple						=       1_000
            _UpperCamelCase      :	Union[str, Any]						=       "huggingface/label-files"
            _UpperCamelCase      :	Optional[Any]						=       "imagenet-1k-id2label.json"
            _UpperCamelCase      :	int						=       json.load(open(hf_hub_download(lowercase_		,lowercase_		,repo_type="dataset"			)		,"r"			)			)
            _UpperCamelCase      :	int						=       {int(lowercase_			): v for k, v in idalabel.items()}
            _UpperCamelCase      :	int						=       idalabel
            _UpperCamelCase      :	List[str]						=       {v: k for k, v in idalabel.items()}
      # size of the architecture
      if model_name in ["dino_vits8", "dino_vits16"]:
            _UpperCamelCase      :	Dict						=       384
            _UpperCamelCase      :	Any						=       1_536
            _UpperCamelCase      :	List[Any]						=       12
            _UpperCamelCase      :	Union[str, Any]						=       6

      # load original model from torch hub
      _UpperCamelCase      :	List[str]						=       torch.hub.load("facebookresearch/dino:main"		,lowercase_			)
      original_model.eval()

      # load state_dict of original model, remove and rename some keys
      _UpperCamelCase      :	int						=       original_model.state_dict()
      if base_model:
            remove_classification_head_(lowercase_			)
      _UpperCamelCase      :	Optional[Any]						=       create_rename_keys(lowercase_		,base_model=lowercase_			)
      for src, dest in rename_keys:
            rename_key(lowercase_		,lowercase_		,lowercase_			)
      read_in_q_k_v(lowercase_		,lowercase_		,lowercase_			)

      # load HuggingFace model
      if base_model:
            _UpperCamelCase      :	Dict						=       ViTModel(lowercase_		,add_pooling_layer=lowercase_			).eval()
      else:
            _UpperCamelCase      :	List[str]						=       ViTForImageClassification(lowercase_			).eval()
      model.load_state_dict(lowercase_			)

      # Check outputs on an image, prepared by ViTImageProcessor
      _UpperCamelCase      :	Tuple						=       ViTImageProcessor()
      _UpperCamelCase      :	str						=       image_processor(images=prepare_img()		,return_tensors="pt"			)
      _UpperCamelCase      :	str						=       encoding["pixel_values"]
      _UpperCamelCase      :	Union[str, Any]						=       model(lowercase_			)

      if base_model:
            _UpperCamelCase      :	Union[str, Any]						=       original_model(lowercase_			)
            assert torch.allclose(lowercase_		,outputs.last_hidden_state[:, 0, :]		,atol=1e-1			)
      else:
            _UpperCamelCase      :	str						=       original_model(lowercase_			)
            assert logits.shape == outputs.logits.shape
            assert torch.allclose(lowercase_		,outputs.logits		,atol=1e-3			)

      Path(lowercase_			).mkdir(exist_ok=lowercase_			)
      print(F'''Saving model {model_name} to {pytorch_dump_folder_path}'''			)
      model.save_pretrained(lowercase_			)
      print(F'''Saving image processor to {pytorch_dump_folder_path}'''			)
      image_processor.save_pretrained(lowercase_			)


if __name__ == "__main__":
  lowerCamelCase__				  =   argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--model_name",
      default="dino_vitb16",
      type=str,
      help="Name of the model trained with DINO you'd like to convert.",
  )
  parser.add_argument(
      "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
  )
  parser.add_argument(
      "--base_model",
      action="store_true",
      help="Whether to only convert the base model (no projection head weights).",
  )

  parser.set_defaults(base_model=True)
  lowerCamelCase__				  =   parser.parse_args()
  convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)







"""simple docstring"""


import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


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


lowerCamelCase__				  =   ["model.decoder.embed_positions.weights"]




def 					lowercase__    (					lowercase_			)					->		Optional[Any]:




      """simple docstring"""

      if "emb" in name:
            _UpperCamelCase      :	List[str]						=       name.replace("emb"		,"model.decoder.embed_tokens"			)
      if "transformer" in name:
            _UpperCamelCase      :	Optional[int]						=       name.replace("transformer"		,"model.decoder"			)
      if "cross_attention" in name:
            _UpperCamelCase      :	Optional[int]						=       name.replace("cross_attention"		,"encoder_attn"			)
      if "linear1" in name:
            _UpperCamelCase      :	Optional[Any]						=       name.replace("linear1"		,"fc1"			)
      if "linear2" in name:
            _UpperCamelCase      :	Union[str, Any]						=       name.replace("linear2"		,"fc2"			)
      if "norm1" in name:
            _UpperCamelCase      :	Optional[Any]						=       name.replace("norm1"		,"self_attn_layer_norm"			)
      if "norm_cross" in name:
            _UpperCamelCase      :	Dict						=       name.replace("norm_cross"		,"encoder_attn_layer_norm"			)
      if "norm2" in name:
            _UpperCamelCase      :	Union[str, Any]						=       name.replace("norm2"		,"final_layer_norm"			)
      if "out_norm" in name:
            _UpperCamelCase      :	Union[str, Any]						=       name.replace("out_norm"		,"model.decoder.layer_norm"			)
      if "linears" in name:
            _UpperCamelCase      :	List[str]						=       name.replace("linears"		,"lm_heads"			)
      if "condition_provider.conditioners.description.output_proj" in name:
            _UpperCamelCase      :	Any						=       name.replace("condition_provider.conditioners.description.output_proj"		,"enc_to_dec_proj"			)
      return name




def 					lowercase__    (					lowercase_		,lowercase_			)					->		Tuple[Dict, Dict]:




      """simple docstring"""

      _UpperCamelCase      :	str						=       list(state_dict.keys()			)
      _UpperCamelCase      :	Optional[Any]						=       {}
      for key in keys:
            _UpperCamelCase      :	Optional[int]						=       state_dict.pop(lowercase_			)
            _UpperCamelCase      :	List[Any]						=       rename_keys(lowercase_			)
            if "in_proj_weight" in key:
                  # split fused qkv proj
                  _UpperCamelCase      :	Tuple						=       val[:hidden_size, :]
                  _UpperCamelCase      :	Optional[Any]						=       val[hidden_size : 2 * hidden_size, :]
                  _UpperCamelCase      :	Optional[Any]						=       val[-hidden_size:, :]
            elif "enc_to_dec_proj" in key:
                  _UpperCamelCase      :	Optional[Any]						=       val
            else:
                  _UpperCamelCase      :	List[str]						=       val
      return state_dict, enc_dec_proj_state_dict




def 					lowercase__    (					lowercase_			)					->		MusicgenDecoderConfig:




      """simple docstring"""

      if checkpoint == "small":
            # default config values
            _UpperCamelCase      :	List[Any]						=       1_024
            _UpperCamelCase      :	List[str]						=       24
            _UpperCamelCase      :	Any						=       16
      elif checkpoint == "medium":
            _UpperCamelCase      :	Tuple						=       1_536
            _UpperCamelCase      :	Dict						=       48
            _UpperCamelCase      :	Tuple						=       24
      elif checkpoint == "large":
            _UpperCamelCase      :	int						=       2_048
            _UpperCamelCase      :	Optional[int]						=       48
            _UpperCamelCase      :	Dict						=       32
      else:
            raise ValueError(F'''Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.'''			)
      _UpperCamelCase      :	str						=       MusicgenDecoderConfig(
          hidden_size=lowercase_		,ffn_dim=hidden_size * 4		,num_hidden_layers=lowercase_		,num_attention_heads=lowercase_		,)
      return config




@torch.no_grad()
def 					lowercase__    (					lowercase_		,lowercase_=None		,lowercase_=None		,lowercase_="cpu"			)					->		List[str]:




      """simple docstring"""

      _UpperCamelCase      :	str						=       MusicGen.get_pretrained(lowercase_		,device=lowercase_			)
      _UpperCamelCase      :	Union[str, Any]						=       decoder_config_from_checkpoint(lowercase_			)

      _UpperCamelCase      :	Optional[int]						=       fairseq_model.lm.state_dict()
      _UpperCamelCase,  _UpperCamelCase      :	Optional[Any]						=       rename_state_dict(
          lowercase_		,hidden_size=decoder_config.hidden_size			)

      _UpperCamelCase      :	Tuple						=       TaEncoderModel.from_pretrained("t5-base"			)
      _UpperCamelCase      :	Union[str, Any]						=       EncodecModel.from_pretrained("facebook/encodec_32khz"			)
      _UpperCamelCase      :	str						=       MusicgenForCausalLM(lowercase_			).eval()

      # load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
      _UpperCamelCase,  _UpperCamelCase      :	str						=       decoder.load_state_dict(lowercase_		,strict=lowercase_			)

      for key in missing_keys.copy():
            if key.startswith(("text_encoder", "audio_encoder")			) or key in EXPECTED_MISSING_KEYS:
                  missing_keys.remove(lowercase_			)

      if len(lowercase_			) > 0:
            raise ValueError(F'''Missing key(s) in state_dict: {missing_keys}'''			)

      if len(lowercase_			) > 0:
            raise ValueError(F'''Unexpected key(s) in state_dict: {unexpected_keys}'''			)

      # init the composite model
      _UpperCamelCase      :	str						=       MusicgenForConditionalGeneration(text_encoder=lowercase_		,audio_encoder=lowercase_		,decoder=lowercase_			)

      # load the pre-trained enc-dec projection (from the decoder state dict)
      model.enc_to_dec_proj.load_state_dict(lowercase_			)

      # check we can do a forward pass
      _UpperCamelCase      :	List[str]						=       torch.arange(0		,8		,dtype=torch.long			).reshape(2		,-1			)
      _UpperCamelCase      :	Dict						=       input_ids.reshape(2 * 4		,-1			)

      with torch.no_grad():
            _UpperCamelCase      :	Tuple						=       model(input_ids=lowercase_		,decoder_input_ids=lowercase_			).logits

      if logits.shape != (8, 1, 2_048):
            raise ValueError("Incorrect shape for logits"			)

      # now construct the processor
      _UpperCamelCase      :	int						=       AutoTokenizer.from_pretrained("t5-base"			)
      _UpperCamelCase      :	str						=       AutoFeatureExtractor.from_pretrained("facebook/encodec_32khz"		,padding_side="left"			)

      _UpperCamelCase      :	Optional[int]						=       MusicgenProcessor(feature_extractor=lowercase_		,tokenizer=lowercase_			)

      # set the appropriate bos/pad token ids
      _UpperCamelCase      :	str						=       2_048
      _UpperCamelCase      :	str						=       2_048

      # set other default generation config params
      _UpperCamelCase      :	Optional[Any]						=       int(30 * audio_encoder.config.frame_rate			)
      _UpperCamelCase      :	List[str]						=       True
      _UpperCamelCase      :	int						=       3.0

      if pytorch_dump_folder is not None:
            Path(lowercase_			).mkdir(exist_ok=lowercase_			)
            logger.info(F'''Saving model {checkpoint} to {pytorch_dump_folder}'''			)
            model.save_pretrained(lowercase_			)
            processor.save_pretrained(lowercase_			)

      if repo_id:
            logger.info(F'''Pushing model {checkpoint} to {repo_id}'''			)
            model.push_to_hub(lowercase_			)
            processor.push_to_hub(lowercase_			)


if __name__ == "__main__":
  lowerCamelCase__				  =   argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--checkpoint",
      default="small",
      type=str,
      help="Checkpoint size of the MusicGen model you'd like to convert. Can be one of: `['small', 'medium', 'large']`.",
  )
  parser.add_argument(
      "--pytorch_dump_folder",
      required=True,
      default=None,
      type=str,
      help="Path to the output PyTorch model directory.",
  )
  parser.add_argument(
      "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
  )
  parser.add_argument(
      "--device", default="cpu", type=str, help="Torch device to run the conversion, either cpu or cuda."
  )

  lowerCamelCase__				  =   parser.parse_args()
  convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)
