
"""simple docstring"""


import argparse
import json
import os

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

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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






def  _snake_case			(						snake_case__				:       List[Any]       ):
	A  				=				ConvNextVaConfig()

	if "atto" in checkpoint_url:
		A  				=				[2, 2, 6, 2]
		A  				=				[40, 80, 160, 320]
	if "femto" in checkpoint_url:
		A  				=				[2, 2, 6, 2]
		A  				=				[48, 96, 192, 384]
	if "pico" in checkpoint_url:
		A  				=				[2, 2, 6, 2]
		A  				=				[64, 128, 256, 512]
	if "nano" in checkpoint_url:
		A  				=				[2, 2, 8, 2]
		A  				=				[80, 160, 320, 640]
	if "tiny" in checkpoint_url:
		A  				=				[3, 3, 9, 3]
		A  				=				[96, 192, 384, 768]
	if "base" in checkpoint_url:
		A  				=				[3, 3, 27, 3]
		A  				=				[128, 256, 512, 1024]
	if "large" in checkpoint_url:
		A  				=				[3, 3, 27, 3]
		A  				=				[192, 384, 768, 1536]
	if "huge" in checkpoint_url:
		A  				=				[3, 3, 27, 3]
		A  				=				[352, 704, 1408, 2816]

	A  				=				1000
	A  				=				"imagenet-1k-id2label.json"
	A  				=				(1, 1000)

	A  				=				"huggingface/label-files"
	A  				=				num_labels
	A  				=				json.load(open(hf_hub_download(_a		, _a		, repo_type='dataset'       )		, 'r'       )       )
	A  				=				{int(_a       ): v for k, v in idalabel.items()}

	A  				=				idalabel
	A  				=				{v: k for k, v in idalabel.items()}
	A  				=				hidden_sizes
	A  				=				depths

	return config, expected_shape






def  _snake_case			(						snake_case__				:       Optional[int]       ):
	if "downsample_layers.0.0" in name:
		A  				=				name.replace('downsample_layers.0.0'		, 'embeddings.patch_embeddings'       )
	if "downsample_layers.0.1" in name:
		A  				=				name.replace('downsample_layers.0.1'		, 'embeddings.norm'       )  # we rename to layernorm later on
	if "downsample_layers.1.0" in name:
		A  				=				name.replace('downsample_layers.1.0'		, 'stages.1.downsampling_layer.0'       )
	if "downsample_layers.1.1" in name:
		A  				=				name.replace('downsample_layers.1.1'		, 'stages.1.downsampling_layer.1'       )
	if "downsample_layers.2.0" in name:
		A  				=				name.replace('downsample_layers.2.0'		, 'stages.2.downsampling_layer.0'       )
	if "downsample_layers.2.1" in name:
		A  				=				name.replace('downsample_layers.2.1'		, 'stages.2.downsampling_layer.1'       )
	if "downsample_layers.3.0" in name:
		A  				=				name.replace('downsample_layers.3.0'		, 'stages.3.downsampling_layer.0'       )
	if "downsample_layers.3.1" in name:
		A  				=				name.replace('downsample_layers.3.1'		, 'stages.3.downsampling_layer.1'       )
	if "stages" in name and "downsampling_layer" not in name:
		# stages.0.0. for instance should be renamed to stages.0.layers.0.
		A  				=				name[: len('stages.0'       )] + ".layers" + name[len('stages.0'       ) :]
	if "gamma" in name:
		A  				=				name.replace('gamma'		, 'weight'       )
	if "beta" in name:
		A  				=				name.replace('beta'		, 'bias'       )
	if "stages" in name:
		A  				=				name.replace('stages'		, 'encoder.stages'       )
	if "norm" in name:
		A  				=				name.replace('norm'		, 'layernorm'       )
	if "head" in name:
		A  				=				name.replace('head'		, 'classifier'       )

	return name






def  _snake_case			(						):
	A  				=				"http://images.cocodataset.org/val2017/000000039769.jpg"
	A  				=				Image.open(requests.get(_a		, stream=_a       ).raw       )
	return im






def  _snake_case			(						snake_case__				:       int       ):
	if "224" in checkpoint_url:
		A  				=				224
		A  				=				224 / 256
	elif "384" in checkpoint_url:
		A  				=				384
		A  				=				None
	else:
		A  				=				512
		A  				=				None

	return ConvNextImageProcessor(
	    size=_a		, crop_pct=_a		, image_mean=[0.485, 0.456, 0.406]		, image_std=[0.229, 0.224, 0.225]		, resample=PILImageResampling.BICUBIC		, )






@torch.no_grad()
def  _snake_case			(						snake_case__				:       Optional[Any]		, snake_case__				:       int		, snake_case__				:       str		, snake_case__				:       List[str]       ):
	print('Downloading original model from checkpoint...'       )
	# define ConvNeXTV2 configuration based on URL
	A  				=				get_convnextva_config(_a       )
	# load original state_dict from URL
	A  				=				torch.hub.load_state_dict_from_url(_a       )["model"]

	print('Converting model parameters...'       )
	# rename keys
	for key in state_dict.copy().keys():
		A  				=				state_dict.pop(_a       )
		A  				=				val
	# add prefix to all keys expect classifier head
	for key in state_dict.copy().keys():
		A  				=				state_dict.pop(_a       )
		if not key.startswith('classifier'       ):
			A  				=				"convnextv2." + key
		A  				=				val

	# load HuggingFace model
	A  				=				ConvNextVaForImageClassification(_a       )
	model.load_state_dict(_a       )
	model.eval()

	# Check outputs on an image, prepared by ConvNextImageProcessor
	A  				=				convert_preprocessor(_a       )
	A  				=				preprocessor(images=prepare_img()		, return_tensors='pt'       )
	A  				=				model(**_a       ).logits

	# note: the logits below were obtained without center cropping
	if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
		A  				=				torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
		A  				=				torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
		A  				=				torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
		A  				=				torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
		A  				=				torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
		A  				=				torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
		A  				=				torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
		A  				=				torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
		A  				=				torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
		A  				=				torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
		A  				=				torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
		A  				=				torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
		A  				=				torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
		A  				=				torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
		A  				=				torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
		A  				=				torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
		A  				=				torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309]       )
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
		A  				=				torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826]       )
	else:
		raise ValueError(F'Unknown URL: {checkpoint_url}'       )

	assert torch.allclose(logits[0, :5]		, _a		, atol=1e-3       )
	assert logits.shape == expected_shape
	print('Model outputs match the original results!'       )

	if save_model:
		print('Saving model to local...'       )
		# Create folder to save model
		if not os.path.isdir(_a       ):
			os.mkdir(_a       )

		model.save_pretrained(_a       )
		preprocessor.save_pretrained(_a       )

	A  				=				"convnextv2"
	if "atto" in checkpoint_url:
		model_name += "-atto"
	if "femto" in checkpoint_url:
		model_name += "-femto"
	if "pico" in checkpoint_url:
		model_name += "-pico"
	if "nano" in checkpoint_url:
		model_name += "-nano"
	elif "tiny" in checkpoint_url:
		model_name += "-tiny"
	elif "base" in checkpoint_url:
		model_name += "-base"
	elif "large" in checkpoint_url:
		model_name += "-large"
	elif "huge" in checkpoint_url:
		model_name += "-huge"
	if "22k" in checkpoint_url and "1k" not in checkpoint_url:
		model_name += "-22k"
	elif "22k" in checkpoint_url and "1k" in checkpoint_url:
		model_name += "-22k-1k"
	elif "1k" in checkpoint_url:
		model_name += "-1k"
	if "224" in checkpoint_url:
		model_name += "-224"
	elif "384" in checkpoint_url:
		model_name += "-384"
	elif "512" in checkpoint_url:
		model_name += "-512"

	if push_to_hub:
		print(F'Pushing {model_name} to the hub...'       )
		model.push_to_hub(_a       )
		preprocessor.push_to_hub(_a       )


if __name__ == "__main__":
		_lowercase											=			argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--checkpoint_url''',
		    default='''https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt''',
		    type=str,
		    help='''URL of the original ConvNeXTV2 checkpoint you\'d like to convert.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''',
		    default='''model''',
		    type=str,
		    help='''Path to the output PyTorch model directory.''',
		)
		parser.add_argument('''--save_model''', action='''store_true''', help='''Save model to local''')
		parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image preprocessor to the hub''')

		_lowercase											=			parser.parse_args()
		convert_convnextva_checkpoint(
		    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
		)

"""simple docstring"""


import argparse
import json
import os.path
from collections import OrderedDict

import numpy as np
import requests
import torch
from flax.training.checkpoints import restore_checkpoint
from huggingface_hub import hf_hub_download

from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
from transformers.image_utils import PILImageResampling






def  _snake_case			(						snake_case__				:       Tuple       ):
	A  				=				'https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint'

	with open(snake_case__		, 'wb'       ) as f:
		with requests.get(snake_case__		, stream=snake_case__       ) as req:
			for chunk in req.iter_content(chunk_size=2048       ):
				f.write(snake_case__       )






def  _snake_case			(						):
	A  				=				VivitConfig()

	A  				=				400
	A  				=				'huggingface/label-files'
	A  				=				'kinetics400-id2label.json'

	A  				=				json.load(open(hf_hub_download(snake_case__		, snake_case__		, repo_type='dataset'       )		, 'r'       )       )
	A  				=				{int(snake_case__       ): v for k, v in idalabel.items()}
	A  				=				idalabel
	A  				=				{v: k for k, v in idalabel.items()}
	return config






def  _snake_case			(						):
	A  				=				hf_hub_download(
	    repo_id='hf-internal-testing/spaghetti-video'		, filename='eating_spaghetti_32_frames.npy'		, repo_type='dataset'       )
	A  				=				np.load(snake_case__       )
	return list(snake_case__       )






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	if np.ndim(snake_case__       ) == 2:
		return transform_attention_bias(snake_case__       )

	elif np.ndim(snake_case__       ) == 3:
		return transform_attention_kernel(snake_case__       )

	else:
		raise Exception(F'Invalid number of dimesions: {np.ndim(snake_case__       )}'       )






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	return current.flatten()






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	return np.reshape(snake_case__		, (current.shape[0], current.shape[1] * current.shape[2])       ).T






def  _snake_case			(						snake_case__				:       np.ndarray       ):
	return np.reshape(snake_case__		, (current.shape[0] * current.shape[1], current.shape[2])       ).T






def  _snake_case			(						snake_case__				:       str		, snake_case__				:       List[str]       ):
	A  				=				state_dict['optimizer']['target']['Transformer'][F'encoderblock_{i}']

	A  				=				F'encoder.layer.{i}.'
	A  				=				{
	    prefix + 'intermediate.dense.bias': state['MlpBlock_0']['Dense_0']['bias'],
	    prefix + 'intermediate.dense.weight': np.transpose(state['MlpBlock_0']['Dense_0']['kernel']       ),
	    prefix + 'output.dense.bias': state['MlpBlock_0']['Dense_1']['bias'],
	    prefix + 'output.dense.weight': np.transpose(state['MlpBlock_0']['Dense_1']['kernel']       ),
	    prefix + 'layernorm_before.bias': state['LayerNorm_0']['bias'],
	    prefix + 'layernorm_before.weight': state['LayerNorm_0']['scale'],
	    prefix + 'layernorm_after.bias': state['LayerNorm_1']['bias'],
	    prefix + 'layernorm_after.weight': state['LayerNorm_1']['scale'],
	    prefix
	    + 'attention.attention.query.bias': transform_attention(
	        state['MultiHeadDotProductAttention_0']['query']['bias']       ),
	    prefix
	    + 'attention.attention.query.weight': transform_attention(
	        state['MultiHeadDotProductAttention_0']['query']['kernel']       ),
	    prefix
	    + 'attention.attention.key.bias': transform_attention(state['MultiHeadDotProductAttention_0']['key']['bias']       ),
	    prefix
	    + 'attention.attention.key.weight': transform_attention(
	        state['MultiHeadDotProductAttention_0']['key']['kernel']       ),
	    prefix
	    + 'attention.attention.value.bias': transform_attention(
	        state['MultiHeadDotProductAttention_0']['value']['bias']       ),
	    prefix
	    + 'attention.attention.value.weight': transform_attention(
	        state['MultiHeadDotProductAttention_0']['value']['kernel']       ),
	    prefix + 'attention.output.dense.bias': state['MultiHeadDotProductAttention_0']['out']['bias'],
	    prefix
	    + 'attention.output.dense.weight': transform_attention_output_weight(
	        state['MultiHeadDotProductAttention_0']['out']['kernel']       ),
	}

	return new_state






def  _snake_case			(						snake_case__				:       Dict       ):
	return sum([1 if 'encoderblock_' in k else 0 for k in state_dict['optimizer']['target']['Transformer'].keys()]       )






def  _snake_case			(						snake_case__				:       Union[str, Any]		, snake_case__				:       Optional[int]=False       ):
	A  				=				get_n_layers(snake_case__       )

	A  				=				OrderedDict()

	A  				=				state_dict['optimizer']['target']['Transformer']['encoder_norm']['bias']
	A  				=				state_dict['optimizer']['target']['Transformer']['encoder_norm']['scale']

	A  				=				np.transpose(
	    state_dict['optimizer']['target']['embedding']['kernel']		, (4, 3, 0, 1, 2)       )
	A  				=				state_dict['optimizer']['target']['embedding']['bias']

	A  				=				state_dict['optimizer']['target']['cls']
	A  				=				state_dict['optimizer']['target']['Transformer']['posembed_input'][
	    'pos_embedding'
	]

	for i in range(snake_case__       ):
		new_state.update(transform_state_encoder_block(snake_case__		, snake_case__       )       )

	if classification_head:
		A  				=				{'vivit.' + k: v for k, v in new_state.items()}
		A  				=				np.transpose(state_dict['optimizer']['target']['output_projection']['kernel']       )
		A  				=				np.transpose(state_dict['optimizer']['target']['output_projection']['bias']       )

	return {k: torch.tensor(snake_case__       ) for k, v in new_state.items()}






def  _snake_case			(						):
	A  				=				VivitImageProcessor()

	assert extractor.do_resize is True
	assert extractor.size == {"shortest_edge": 256}
	assert extractor.do_center_crop is True
	assert extractor.crop_size == {"width": 224, "height": 224}
	assert extractor.resample == PILImageResampling.BILINEAR

	# here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
	# one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
	# which effectively means no normalization (and ViViT does not overwrite those when calling this func)
	assert extractor.do_normalize is False
	assert extractor.do_rescale is True
	assert extractor.rescale_factor == 1 / 255

	# zero-centering = True in original implementation
	assert extractor.do_zero_centering is True

	return extractor






def  _snake_case			(						snake_case__				:       str       ):
	A  				=				'checkpoint'

	if not os.path.exists(snake_case__       ):
		download_checkpoint(snake_case__       )

	A  				=				restore_checkpoint(snake_case__		, snake_case__       )
	A  				=				transform_state(snake_case__		, classification_head=snake_case__       )

	A  				=				get_vivit_config()

	assert config.image_size == 224
	assert config.num_frames == 32

	A  				=				VivitForVideoClassification(snake_case__       )
	model.load_state_dict(snake_case__       )
	model.eval()

	A  				=				get_processor()

	A  				=				prepare_video()
	A  				=				extractor(snake_case__		, return_tensors='pt'       )

	A  				=				model(**snake_case__       )

	A  				=				torch.Size([1, 400]       )
	A  				=				torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658]       )

	assert outputs.logits.shape == expected_shape
	assert torch.allclose(outputs.logits[0, :5]		, snake_case__		, atol=1e-4       ), outputs.logits[0, :5]

	model.save_pretrained(snake_case__       )
	extractor.save_pretrained(snake_case__       )


if __name__ == "__main__":
		_lowercase											=			argparse.ArgumentParser()

		parser.add_argument('''--output_model_name''', '''-o''', type=str, help='''Output path for the converted HuggingFace model''')

		_lowercase											=			parser.parse_args()
		convert(args.output_model_name)