

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







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase  )    ->    int:
	UpperCAmelCase_		:      Any       	=					['head.weight', 'head.bias']
	for k in ignore_keys:
		state_dict.pop(lowerCamelCase__   ,lowerCamelCase__  )







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase  )    ->    Optional[int]:
	# projection head is used in the self-supervised pre-training in MSN,
	# for downstream task it's not needed.
	UpperCAmelCase_		:      Dict       	=					[
	    '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(lowerCamelCase__   ,lowerCamelCase__  )







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase  )    ->    Optional[int]:
	UpperCAmelCase_		:      int       	=					dct.pop(lowerCamelCase__  )
	UpperCAmelCase_		:      str       	=					val







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase   ,UpperCamelCase  )    ->    Optional[Any]:
	UpperCAmelCase_		:      Union[str, Any]       	=					ViTMSNConfig()
	UpperCAmelCase_		:      Optional[Any]       	=					1_0_0_0

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

	if "s16" in checkpoint_url:
		UpperCAmelCase_		:      List[str]       	=					3_8_4
		UpperCAmelCase_		:      List[str]       	=					1_5_3_6
		UpperCAmelCase_		:      Dict       	=					6
	elif "l16" in checkpoint_url:
		UpperCAmelCase_		:      Any       	=					1_0_2_4
		UpperCAmelCase_		:      Tuple       	=					4_0_9_6
		UpperCAmelCase_		:      int       	=					2_4
		UpperCAmelCase_		:      List[str]       	=					1_6
		UpperCAmelCase_		:      int       	=					0.1
	elif "b4" in checkpoint_url:
		UpperCAmelCase_		:      List[Any]       	=					4
	elif "l7" in checkpoint_url:
		UpperCAmelCase_		:      Optional[Any]       	=					7
		UpperCAmelCase_		:      Optional[Any]       	=					1_0_2_4
		UpperCAmelCase_		:      Tuple       	=					4_0_9_6
		UpperCAmelCase_		:      List[Any]       	=					2_4
		UpperCAmelCase_		:      Optional[int]       	=					1_6
		UpperCAmelCase_		:      str       	=					0.1

	UpperCAmelCase_		:      List[Any]       	=					ViTMSNModel(lowerCamelCase__  )

	UpperCAmelCase_		:      Optional[int]       	=					torch.hub.load_state_dict_from_url(lowerCamelCase__   ,map_location='cpu'  )['target_encoder']

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

	remove_projection_head(lowerCamelCase__  )
	UpperCAmelCase_		:      Optional[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__   ,base_model=lowerCamelCase__  )

	model.load_state_dict(lowerCamelCase__  )
	model.eval()

	UpperCAmelCase_		:      List[str]       	=					'http://images.cocodataset.org/val2017/000000039769.jpg'

	UpperCAmelCase_		:      Tuple       	=					Image.open(requests.get(lowerCamelCase__   ,stream=lowerCamelCase__  ).raw  )
	UpperCAmelCase_		:      Union[str, Any]       	=					ViTImageProcessor(
	    size=config.image_size   ,image_mean=lowerCamelCase__   ,image_std=lowerCamelCase__  )
	UpperCAmelCase_		:      List[Any]       	=					image_processor(images=lowerCamelCase__   ,return_tensors='pt'  )

	# forward pass
	torch.manual_seed(2  )
	UpperCAmelCase_		:      Any       	=					model(**lowerCamelCase__  )
	UpperCAmelCase_		:      Union[str, Any]       	=					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:
		UpperCAmelCase_		:      List[Any]       	=					torch.tensor([[-1.0915, -1.4876, -1.1809]]  )
	elif "b16" in checkpoint_url:
		UpperCAmelCase_		:      Union[str, Any]       	=					torch.tensor([[14.2889, -18.9045, 11.7281]]  )
	elif "l16" in checkpoint_url:
		UpperCAmelCase_		:      int       	=					torch.tensor([[41.5028, -22.8681, 45.6475]]  )
	elif "b4" in checkpoint_url:
		UpperCAmelCase_		:      List[Any]       	=					torch.tensor([[-4.3868, 5.2932, -0.4137]]  )
	else:
		UpperCAmelCase_		:      Optional[Any]       	=					torch.tensor([[-0.1792, -0.6465, 2.4263]]  )

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

	print(f"""Saving model 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(
						    "--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."
						)

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









'''simple docstring'''


import importlib.metadata
import warnings
from copy import deepcopy

from packaging import version

from ..utils import logging
from .import_utils import is_accelerate_available, is_bitsandbytes_available


if is_bitsandbytes_available():
						import bitsandbytes as bnb
						import torch
						import torch.nn as nn

						from ..pytorch_utils import ConvaD

if is_accelerate_available():
						from accelerate import init_empty_weights
						from accelerate.utils import find_tied_parameters

lowerCAmelCase__										= logging.get_logger(__name__)







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase=None   ,UpperCamelCase=None  )    ->    Any:
	# Recurse if needed
	if "." in tensor_name:
		UpperCAmelCase_		:      Tuple       	=					tensor_name.split('.'  )
		for split in splits[:-1]:
			UpperCAmelCase_		:      Optional[int]       	=					getattr(UpperCamelCase   ,UpperCamelCase  )
			if new_module is None:
				raise ValueError(f"""{module} has no attribute {split}."""  )
			UpperCAmelCase_		:      Union[str, Any]       	=					new_module
		UpperCAmelCase_		:      int       	=					splits[-1]

	if tensor_name not in module._parameters and tensor_name not in module._buffers:
		raise ValueError(f"""{module} does not have a parameter or a buffer named {tensor_name}."""  )
	UpperCAmelCase_		:      Tuple       	=					tensor_name in module._buffers
	UpperCAmelCase_		:      Any       	=					getattr(UpperCamelCase   ,UpperCamelCase  )

	if old_value.device == torch.device('meta'  ) and device not in ["meta", torch.device('meta'  )] and value is None:
		raise ValueError(f"""{tensor_name} is on the meta device, we need a `value` to put in on {device}."""  )

	UpperCAmelCase_		:      Any       	=					False
	UpperCAmelCase_		:      Any       	=					False
	if is_buffer or not is_bitsandbytes_available():
		UpperCAmelCase_		:      Dict       	=					False
		UpperCAmelCase_		:      Tuple       	=					False
	else:
		UpperCAmelCase_		:      List[str]       	=					hasattr(bnb.nn   ,'Params4bit'  ) and isinstance(module._parameters[tensor_name]   ,bnb.nn.Paramsabit  )
		UpperCAmelCase_		:      str       	=					isinstance(module._parameters[tensor_name]   ,bnb.nn.IntaParams  )

	if is_abit or is_abit:
		UpperCAmelCase_		:      Union[str, Any]       	=					module._parameters[tensor_name]
		if param.device.type != "cuda":
			if value is None:
				UpperCAmelCase_		:      List[str]       	=					old_value.to(UpperCamelCase  )
			elif isinstance(UpperCamelCase   ,torch.Tensor  ):
				UpperCAmelCase_		:      Any       	=					value.to('cpu'  )
				if value.dtype == torch.inta:
					UpperCAmelCase_		:      Dict       	=					version.parse(importlib.metadata.version('bitsandbytes'  )  ) > version.parse(
					    '0.37.2'  )
					if not is_abit_serializable:
						raise ValueError(
						    'Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. '
						    'Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`.'  )
			else:
				UpperCAmelCase_		:      List[Any]       	=					torch.tensor(UpperCamelCase   ,device='cpu'  )

			# Support models using `Conv1D` in place of `nn.Linear` (e.g. gpt2) by transposing the weight matrix prior to quantization.
			# Since weights are saved in the correct "orientation", we skip transposing when loading.
			if issubclass(module.source_cls   ,UpperCamelCase  ) and fpaa_statistics is None:
				UpperCAmelCase_		:      Dict       	=					new_value.T

			UpperCAmelCase_		:      Tuple       	=					old_value.__dict__
			if is_abit:
				UpperCAmelCase_		:      Any       	=					bnb.nn.IntaParams(UpperCamelCase   ,requires_grad=UpperCamelCase   ,**UpperCamelCase  ).to(UpperCamelCase  )
			elif is_abit:
				UpperCAmelCase_		:      Optional[Any]       	=					bnb.nn.Paramsabit(UpperCamelCase   ,requires_grad=UpperCamelCase   ,**UpperCamelCase  ).to(UpperCamelCase  )

			UpperCAmelCase_		:      Optional[Any]       	=					new_value
			if fpaa_statistics is not None:
				setattr(module.weight   ,'SCB'   ,fpaa_statistics.to(UpperCamelCase  )  )

	else:
		if value is None:
			UpperCAmelCase_		:      str       	=					old_value.to(UpperCamelCase  )
		elif isinstance(UpperCamelCase   ,torch.Tensor  ):
			UpperCAmelCase_		:      Tuple       	=					value.to(UpperCamelCase  )
		else:
			UpperCAmelCase_		:      str       	=					torch.tensor(UpperCamelCase   ,device=UpperCamelCase  )

		if is_buffer:
			UpperCAmelCase_		:      Dict       	=					new_value
		else:
			UpperCAmelCase_		:      Optional[int]       	=					nn.Parameter(UpperCamelCase   ,requires_grad=old_value.requires_grad  )
			UpperCAmelCase_		:      Optional[int]       	=					new_value







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase   ,UpperCamelCase=None   ,UpperCamelCase=None   ,UpperCamelCase=None   ,UpperCamelCase=False  )    ->    Tuple:
	for name, module in model.named_children():
		if current_key_name is None:
			UpperCAmelCase_		:      Optional[int]       	=					[]
		current_key_name.append(UpperCamelCase  )

		if (isinstance(UpperCamelCase   ,nn.Linear  ) or isinstance(UpperCamelCase   ,UpperCamelCase  )) and name not in modules_to_not_convert:
			# Check if the current key is not in the `modules_to_not_convert`
			if not any(key in '.'.join(UpperCamelCase  ) for key in modules_to_not_convert  ):
				with init_empty_weights():
					if isinstance(UpperCamelCase   ,UpperCamelCase  ):
						UpperCAmelCase_	,    UpperCAmelCase_		:      Optional[Any]       	=					module.weight.shape
					else:
						UpperCAmelCase_		:      Dict       	=					module.in_features
						UpperCAmelCase_		:      Any       	=					module.out_features

					if quantization_config.quantization_method() == "llm_int8":
						UpperCAmelCase_		:      Optional[int]       	=					bnb.nn.LinearabitLt(
						    UpperCamelCase   ,UpperCamelCase   ,module.bias is not None   ,has_fpaa_weights=quantization_config.llm_inta_has_fpaa_weight   ,threshold=quantization_config.llm_inta_threshold   ,)
						UpperCAmelCase_		:      List[Any]       	=					True
					else:
						if (
						    quantization_config.llm_inta_skip_modules is not None
						    and name in quantization_config.llm_inta_skip_modules
						):
							pass
						else:
							UpperCAmelCase_		:      Union[str, Any]       	=					bnb.nn.Linearabit(
							    UpperCamelCase   ,UpperCamelCase   ,module.bias is not None   ,quantization_config.bnb_abit_compute_dtype   ,compress_statistics=quantization_config.bnb_abit_use_double_quant   ,quant_type=quantization_config.bnb_abit_quant_type   ,)
							UpperCAmelCase_		:      Dict       	=					True
                    # Store the module class in case we need to transpose the weight later
					UpperCAmelCase_		:      List[str]       	=					type(UpperCamelCase  )
					# Force requires grad to False to avoid unexpected errors
					model._modules[name].requires_grad_(UpperCamelCase  )
		if len(list(module.children()  )  ) > 0:
			UpperCAmelCase_	,    UpperCAmelCase_		:      Any       	=					_replace_with_bnb_linear(
			    UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,has_been_replaced=UpperCamelCase   ,)
		# Remove the last key for recursion
		current_key_name.pop(-1  )
	return model, has_been_replaced







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase   ,UpperCamelCase=None   ,UpperCamelCase=None   ,UpperCamelCase=None  )    ->    Optional[int]:
	UpperCAmelCase_		:      Union[str, Any]       	=					['lm_head'] if modules_to_not_convert is None else modules_to_not_convert
	UpperCAmelCase_	,    UpperCAmelCase_		:      Any       	=					_replace_with_bnb_linear(
	    UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase  )

	if not has_been_replaced:
		logger.warning(
		    'You are loading your model in 8bit or 4bit but no linear modules were found in your model.'
		    ' Please double check your model architecture, or submit an issue on github if you think this is'
		    ' a bug.'  )

	return model







def 						SCREAMING_SNAKE_CASE(  *UpperCamelCase   ,**UpperCamelCase  )    ->    Tuple:
	warnings.warn(
	    '`replace_8bit_linear` will be deprecated in a future version, please use `replace_with_bnb_linear` instead'   ,UpperCamelCase   ,)
	return replace_with_bnb_linear(*UpperCamelCase   ,**UpperCamelCase  )







def 						SCREAMING_SNAKE_CASE(  *UpperCamelCase   ,**UpperCamelCase  )    ->    List[Any]:
	warnings.warn(
	    '`set_module_8bit_tensor_to_device` will be deprecated in a future version, please use `set_module_quantized_tensor_to_device` instead'   ,UpperCamelCase   ,)
	return set_module_quantized_tensor_to_device(*UpperCamelCase   ,**UpperCamelCase  )







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase  )    ->    int:
	UpperCAmelCase_		:      Union[str, Any]       	=					deepcopy(UpperCamelCase  )  # this has 0 cost since it is done inside `init_empty_weights` context manager`
	tied_model.tie_weights()

	UpperCAmelCase_		:      Tuple       	=					find_tied_parameters(UpperCamelCase  )
	# For compatibility with Accelerate < 0.18
	if isinstance(UpperCamelCase   ,UpperCamelCase  ):
		UpperCAmelCase_		:      Dict       	=					sum(list(tied_params.values()  )   ,[]  ) + list(tied_params.keys()  )
	else:
		UpperCAmelCase_		:      Union[str, Any]       	=					sum(UpperCamelCase   ,[]  )
	UpperCAmelCase_		:      Union[str, Any]       	=					len(UpperCamelCase  ) > 0

	# Check if it is a base model
	UpperCAmelCase_		:      List[Any]       	=					not hasattr(UpperCamelCase   ,model.base_model_prefix  )

	# Ignore this for base models (BertModel, GPT2Model, etc.)
	if (not has_tied_params) and is_base_model:
		return []

	# otherwise they have an attached head
	UpperCAmelCase_		:      str       	=					list(model.named_children()  )
	UpperCAmelCase_		:      Optional[int]       	=					[list_modules[-1][0]]

	# add last module together with tied weights
	UpperCAmelCase_		:      Any       	=					set(UpperCamelCase  ) - set(UpperCamelCase  )
	UpperCAmelCase_		:      Optional[int]       	=					list(set(UpperCamelCase  )  ) + list(UpperCamelCase  )

	# remove ".weight" from the keys
	UpperCAmelCase_		:      Optional[int]       	=					['.weight', '.bias']
	UpperCAmelCase_		:      Tuple       	=					[]
	for name in list_untouched:
		for name_to_remove in names_to_remove:
			if name_to_remove in name:
				UpperCAmelCase_		:      Union[str, Any]       	=					name.replace(UpperCamelCase   ,''  )
		filtered_module_names.append(UpperCamelCase  )

	return filtered_module_names






