



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 	_lowerCAmelCase						(				__lowerCamelCase   :     Dict		,	__lowerCamelCase   :     Any=False     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Tuple       =	[]
 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"
  __SCREAMING_SNAKE_CASE       :    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 	_lowerCAmelCase						(				__lowerCamelCase   :     Union[str, Any]		,	__lowerCamelCase   :     List[str]		,	__lowerCamelCase   :     Optional[Any]=False     ):



 """simple docstring"""





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






def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[str]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Dict       =	["head.weight", "head.bias"]
 for k in ignore_keys:
  state_dict.pop(__lowerCamelCase		,	__lowerCamelCase     )






def 	_lowerCAmelCase						(				__lowerCamelCase   :     str		,	__lowerCamelCase   :     List[Any]		,	__lowerCamelCase   :     Optional[int]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	dct.pop(__lowerCamelCase     )
 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	val






def 	_lowerCAmelCase						(				):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[Any]       =	"http://images.cocodataset.org/val2017/000000039769.jpg"
 __SCREAMING_SNAKE_CASE       :    Tuple       =	Image.open(requests.get(__lowerCamelCase		,	stream=__lowerCamelCase     ).raw     )
 return im






@torch.no_grad()
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Any		,	__lowerCamelCase   :     int		,	__lowerCamelCase   :     Optional[Any]=True     ):



 """simple docstring"""





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

 # load original model from torch hub
 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	torch.hub.load("facebookresearch/dino:main"		,	__lowerCamelCase     )
 original_model.eval()

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

 # load HuggingFace model
 if base_model:
  __SCREAMING_SNAKE_CASE       :    str       =	ViTModel(__lowerCamelCase		,	add_pooling_layer=__lowerCamelCase     ).eval()
 else:
  __SCREAMING_SNAKE_CASE       :    int       =	ViTForImageClassification(__lowerCamelCase     ).eval()
 model.load_state_dict(__lowerCamelCase     )

 # Check outputs on an image, prepared by ViTImageProcessor
 __SCREAMING_SNAKE_CASE       :    List[Any]       =	ViTImageProcessor()
 __SCREAMING_SNAKE_CASE       :    int       =	image_processor(images=prepare_img()		,	return_tensors="pt"     )
 __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	encoding["pixel_values"]
 __SCREAMING_SNAKE_CASE       :    Dict       =	model(__lowerCamelCase     )

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

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


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)





from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
      from .deepspeed import (
          DeepSpeedEngineWrapper,
          DeepSpeedOptimizerWrapper,
          DeepSpeedSchedulerWrapper,
          DummyOptim,
          DummyScheduler,
          HfDeepSpeedConfig,
      )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers






