'''simple docstring'''


import argparse
import json

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

from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling





def       lowerCAmelCase					(		UpperCamelCase__				:   Optional[int]       ):


	"""simple docstring"""
	__UpperCAmelCase							=			[2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 1_8, 2]
	__UpperCAmelCase							=			True if '''large''' in model_name or '''huge''' in model_name else False
	__UpperCAmelCase							=			True if '''large''' in model_name or '''huge''' in model_name else False
	__UpperCAmelCase							=			True if '''large''' in model_name or '''huge''' in model_name else False

	if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
		if "fl3" in model_name:
			__UpperCAmelCase							=			[3, 3, 3, 3]
			__UpperCAmelCase							=			[5, 5, 5, 5]
		elif "fl4" in model_name:
			__UpperCAmelCase							=			[4, 4, 4, 4]
			__UpperCAmelCase							=			[3, 3, 3, 3]

	if "tiny" in model_name or "small" in model_name or "base" in model_name:
		__UpperCAmelCase							=			[3, 3, 3, 3]
		if "lrf" in model_name:
			__UpperCAmelCase							=			[3, 3, 3, 3]
		else:
			__UpperCAmelCase							=			[2, 2, 2, 2]

	if "tiny" in model_name:
		__UpperCAmelCase							=			9_6
	elif "small" in model_name:
		__UpperCAmelCase							=			9_6
	elif "base" in model_name:
		__UpperCAmelCase							=			1_2_8
	elif "large" in model_name:
		__UpperCAmelCase							=			1_9_2
	elif "xlarge" in model_name:
		__UpperCAmelCase							=			2_5_6
	elif "huge" in model_name:
		__UpperCAmelCase							=			3_5_2

	# set label information
	__UpperCAmelCase							=			'''huggingface/label-files'''
	if "large" in model_name or "huge" in model_name:
		__UpperCAmelCase							=			'''imagenet-22k-id2label.json'''
	else:
		__UpperCAmelCase							=			'''imagenet-1k-id2label.json'''

	__UpperCAmelCase							=			json.load(open(hf_hub_download(UpperCamelCase__    ,					UpperCamelCase__    ,					repo_type='''dataset'''       )    ,					'''r'''       )       )
	__UpperCAmelCase							=			{int(UpperCamelCase__       ): v for k, v in idalabel.items()}
	__UpperCAmelCase							=			{v: k for k, v in idalabel.items()}

	__UpperCAmelCase							=			FocalNetConfig(
	    embed_dim=UpperCamelCase__    ,					depths=UpperCamelCase__    ,					focal_levels=UpperCamelCase__    ,					focal_windows=UpperCamelCase__    ,					use_conv_embed=UpperCamelCase__    ,					idalabel=UpperCamelCase__    ,					labelaid=UpperCamelCase__    ,					use_post_layernorm=UpperCamelCase__    ,					use_layerscale=UpperCamelCase__    ,					)

	return config





def       lowerCAmelCase					(		UpperCamelCase__				:   int       ):


	"""simple docstring"""
	if "patch_embed.proj" in name:
		__UpperCAmelCase							=			name.replace('''patch_embed.proj'''    ,					'''embeddings.patch_embeddings.projection'''       )
	if "patch_embed.norm" in name:
		__UpperCAmelCase							=			name.replace('''patch_embed.norm'''    ,					'''embeddings.norm'''       )
	if "layers" in name:
		__UpperCAmelCase							=			'''encoder.''' + name
	if "encoder.layers" in name:
		__UpperCAmelCase							=			name.replace('''encoder.layers'''    ,					'''encoder.stages'''       )
	if "downsample.proj" in name:
		__UpperCAmelCase							=			name.replace('''downsample.proj'''    ,					'''downsample.projection'''       )
	if "blocks" in name:
		__UpperCAmelCase							=			name.replace('''blocks'''    ,					'''layers'''       )
	if "modulation.f.weight" in name or "modulation.f.bias" in name:
		__UpperCAmelCase							=			name.replace('''modulation.f'''    ,					'''modulation.projection_in'''       )
	if "modulation.h.weight" in name or "modulation.h.bias" in name:
		__UpperCAmelCase							=			name.replace('''modulation.h'''    ,					'''modulation.projection_context'''       )
	if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
		__UpperCAmelCase							=			name.replace('''modulation.proj'''    ,					'''modulation.projection_out'''       )

	if name == "norm.weight":
		__UpperCAmelCase							=			'''layernorm.weight'''
	if name == "norm.bias":
		__UpperCAmelCase							=			'''layernorm.bias'''

	if "head" in name:
		__UpperCAmelCase							=			name.replace('''head'''    ,					'''classifier'''       )
	else:
		__UpperCAmelCase							=			'''focalnet.''' + name

	return name





def       lowerCAmelCase					(		UpperCamelCase__				:   Tuple    ,					UpperCamelCase__				:   Dict    ,					UpperCamelCase__				:   Tuple=False       ):


	"""simple docstring"""
	# fmt: off
	__UpperCAmelCase							=			{
	    '''focalnet-tiny''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth''',
	    '''focalnet-tiny-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth''',
	    '''focalnet-small''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth''',
	    '''focalnet-small-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth''',
	    '''focalnet-base''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth''',
	    '''focalnet-base-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth''',
	    '''focalnet-large-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth''',
	    '''focalnet-large-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth''',
	    '''focalnet-xlarge-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth''',
	    '''focalnet-xlarge-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth''',
	}
	# fmt: on

	__UpperCAmelCase							=			model_name_to_url[model_name]
	print('''Checkpoint URL: '''    ,					UpperCamelCase__       )
	__UpperCAmelCase							=			torch.hub.load_state_dict_from_url(UpperCamelCase__    ,					map_location='''cpu'''       )['''model''']

	# rename keys
	for key in state_dict.copy().keys():
		__UpperCAmelCase							=			state_dict.pop(UpperCamelCase__       )
		__UpperCAmelCase							=			val

	__UpperCAmelCase							=			get_focalnet_config(UpperCamelCase__       )
	__UpperCAmelCase							=			FocalNetForImageClassification(UpperCamelCase__       )
	model.eval()

	# load state dict
	model.load_state_dict(UpperCamelCase__       )

	# verify conversion
	__UpperCAmelCase							=			'''http://images.cocodataset.org/val2017/000000039769.jpg'''

	__UpperCAmelCase							=			BitImageProcessor(
	    do_resize=UpperCamelCase__    ,					size={'''shortest_edge''': 2_5_6}    ,					resample=PILImageResampling.BILINEAR    ,					do_center_crop=UpperCamelCase__    ,					crop_size=2_2_4    ,					do_normalize=UpperCamelCase__    ,					image_mean=UpperCamelCase__    ,					image_std=UpperCamelCase__    ,					)
	__UpperCAmelCase							=			Image.open(requests.get(UpperCamelCase__    ,					stream=UpperCamelCase__       ).raw       )
	__UpperCAmelCase							=			processor(images=UpperCamelCase__    ,					return_tensors='''pt'''       )

	__UpperCAmelCase							=			transforms.Compose(
	    [
	        transforms.Resize(2_5_6       ),
	        transforms.CenterCrop(2_2_4       ),
	        transforms.ToTensor(),
	        transforms.Normalize(mean=[0.4_85, 0.4_56, 0.4_06]    ,					std=[0.2_29, 0.2_24, 0.2_25]       ),
	    ]       )

	__UpperCAmelCase							=			image_transforms(UpperCamelCase__       ).unsqueeze(0       )

	# verify pixel_values
	assert torch.allclose(inputs.pixel_values    ,					UpperCamelCase__    ,					atol=1E-4       )

	__UpperCAmelCase							=			model(**UpperCamelCase__       )

	__UpperCAmelCase							=			outputs.logits.argmax(-1       ).item()
	print('''Predicted class:'''    ,					model.config.idalabel[predicted_class_idx]       )

	print('''First values of logits:'''    ,					outputs.logits[0, :3]       )

	if model_name == "focalnet-tiny":
		__UpperCAmelCase							=			torch.tensor([0.21_66, -0.43_68, 0.21_91]       )
	elif model_name == "focalnet-tiny-lrf":
		__UpperCAmelCase							=			torch.tensor([1.16_69, 0.01_25, -0.16_95]       )
	elif model_name == "focalnet-small":
		__UpperCAmelCase							=			torch.tensor([0.49_17, -0.04_30, 0.13_41]       )
	elif model_name == "focalnet-small-lrf":
		__UpperCAmelCase							=			torch.tensor([-0.25_88, -0.53_42, -0.23_31]       )
	elif model_name == "focalnet-base":
		__UpperCAmelCase							=			torch.tensor([-0.16_55, -0.40_90, -0.17_30]       )
	elif model_name == "focalnet-base-lrf":
		__UpperCAmelCase							=			torch.tensor([0.53_06, -0.04_83, -0.39_28]       )
	assert torch.allclose(outputs.logits[0, :3]    ,					UpperCamelCase__    ,					atol=1E-4       )
	print('''Looks ok!'''       )

	if pytorch_dump_folder_path is not None:
		print(f"""Saving model and processor of {model_name} to {pytorch_dump_folder_path}"""       )
		model.save_pretrained(UpperCamelCase__       )
		processor.save_pretrained(UpperCamelCase__       )

	if push_to_hub:
		print(f"""Pushing model and processor of {model_name} to the hub..."""       )
		model.push_to_hub(f"""{model_name}"""       )
		processor.push_to_hub(f"""{model_name}"""       )


if __name__ == "__main__":
					__lowerCAmelCase	:       Any						       =	argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    "--model_name",
					    default="focalnet-tiny",
					    type=str,
					    help="Name of the FocalNet model 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(
					    "--push_to_hub",
					    action="store_true",
					    help="Whether to push the model and processor to the hub.",
					)

					__lowerCAmelCase	:       Any						       =	parser.parse_args()
					convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)

'''simple docstring'''


import ast
import os
import re
import shutil
import tempfile
import unittest
from unittest import mock

import torch

from accelerate.test_utils.examples import compare_against_test
from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
from accelerate.utils import write_basic_config


# DataLoaders built from `test_samples/MRPC` for quick testing
# Should mock `{script_name}.get_dataloaders` via:
# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)

__lowerCAmelCase	:       Optional[int]						       =	[
    "cross_validation.py",
    "gradient_accumulation.py",
    "local_sgd.py",
    "multi_process_metrics.py",
    "memory.py",
    "automatic_gradient_accumulation.py",
    "fsdp_with_peak_mem_tracking.py",
    "deepspeed_with_config_support.py",
    "megatron_lm_gpt_pretraining.py",
]






class 						A   (   unittest.TestCase   ):





							def       snake_case__						(      self				:							Any     ,		__a				:							str     ,		__a				:							bool     ,		__a				:							str = None     ,		__a				:							list = None	)							->	Tuple:
								__UpperCAmelCase							=			None
								__UpperCAmelCase							=			os.path.abspath(os.path.join('''examples'''     ,		'''by_feature'''	)	)
								__UpperCAmelCase							=			os.path.abspath('''examples'''	)
								for item in os.listdir(__a	):
									if item not in EXCLUDE_EXAMPLES:
										__UpperCAmelCase							=			os.path.join(__a     ,		__a	)
										if os.path.isfile(__a	) and ".py" in item_path:
											with self.subTest(
											    tested_script=__a     ,		feature_script=__a     ,		tested_section='''main()''' if parser_only else '''training_function()'''     ,		):
												__UpperCAmelCase							=			compare_against_test(
												    os.path.join(__a     ,		__a	)     ,		__a     ,		__a     ,		__a	)
												__UpperCAmelCase							=			'''\n'''.join(__a	)
												if special_strings is not None:
													for string in special_strings:
														__UpperCAmelCase							=			diff.replace(__a     ,		''''''	)
												self.assertEqual(__a     ,		''''''	)





							def       snake_case__						(      self				:							Optional[Any]	)							->	str:
								self.one_complete_example('''complete_nlp_example.py'''     ,		__a	)
								self.one_complete_example('''complete_nlp_example.py'''     ,		__a	)






							def       snake_case__						(      self				:							List[str]	)							->	Tuple:
								__UpperCAmelCase							=			os.path.abspath(os.path.join('''examples'''     ,		'''cv_example.py'''	)	)
								__UpperCAmelCase							=			[
								    ''' ''' * 1_6 + '''{\n\n''',
								    ''' ''' * 2_0 + '''"accuracy": eval_metric["accuracy"],\n\n''',
								    ''' ''' * 2_0 + '''"f1": eval_metric["f1"],\n\n''',
								    ''' ''' * 2_0 + '''"train_loss": total_loss.item() / len(train_dataloader),\n\n''',
								    ''' ''' * 2_0 + '''"epoch": epoch,\n\n''',
								    ''' ''' * 1_6 + '''},\n\n''',
								    ''' ''' * 1_6 + '''step=epoch,\n''',
								    ''' ''' * 1_2,
								    ''' ''' * 8 + '''for step, batch in enumerate(active_dataloader):\n''',
								]
								self.one_complete_example('''complete_cv_example.py'''     ,		__a     ,		__a     ,		__a	)
								self.one_complete_example('''complete_cv_example.py'''     ,		__a     ,		__a     ,		__a	)






@mock.patch.dict(os.environ       ,	{'''TESTING_MOCKED_DATALOADERS''': '''1'''}   )
class 						A   (   UpperCAmelCase   ):
							a_        =      False





							@classmethod
							def       snake_case__						(      cls				:							Tuple	)							->	str:
								super().setUpClass()
								__UpperCAmelCase							=			tempfile.mkdtemp()
								__UpperCAmelCase							=			os.path.join(cls._tmpdir     ,		'''default_config.yml'''	)

								write_basic_config(save_location=cls.configPath	)
								__UpperCAmelCase							=			['''accelerate''', '''launch''', '''--config_file''', cls.configPath]





							@classmethod
							def       snake_case__						(      cls				:							Dict	)							->	int:
								super().tearDownClass()
								shutil.rmtree(cls._tmpdir	)





							def       snake_case__						(      self				:							Tuple	)							->	Dict:
								__UpperCAmelCase							=			f"""
        examples/by_feature/checkpointing.py
        --checkpointing_steps epoch
        --output_dir {self.tmpdir}
        """.split()
								run_command(self._launch_args + testargs	)
								self.assertTrue(os.path.exists(os.path.join(self.tmpdir     ,		'''epoch_0'''	)	)	)





							def       snake_case__						(      self				:							str	)							->	int:
								__UpperCAmelCase							=			f"""
        examples/by_feature/checkpointing.py
        --checkpointing_steps 1
        --output_dir {self.tmpdir}
        """.split()
								__UpperCAmelCase							=			run_command(self._launch_args + testargs	)
								self.assertTrue(os.path.exists(os.path.join(self.tmpdir     ,		'''step_2'''	)	)	)





							def       snake_case__						(      self				:							Any	)							->	Any:
								__UpperCAmelCase							=			f"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir     ,		'epoch_0'	)}
        """.split()
								__UpperCAmelCase							=			run_command(self._launch_args + testargs     ,		return_stdout=__a	)
								self.assertNotIn('''epoch 0:'''     ,		__a	)
								self.assertIn('''epoch 1:'''     ,		__a	)





							def       snake_case__						(      self				:							Tuple	)							->	Optional[int]:
								__UpperCAmelCase							=			f"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir     ,		'step_2'	)}
        """.split()
								__UpperCAmelCase							=			run_command(self._launch_args + testargs     ,		return_stdout=__a	)
								if torch.cuda.is_available():
									__UpperCAmelCase							=			torch.cuda.device_count()
								else:
									__UpperCAmelCase							=			1
								if num_processes > 1:
									self.assertNotIn('''epoch 0:'''     ,		__a	)
									self.assertIn('''epoch 1:'''     ,		__a	)
								else:
									self.assertIn('''epoch 0:'''     ,		__a	)
									self.assertIn('''epoch 1:'''     ,		__a	)





							@slow
							def       snake_case__						(      self				:							Any	)							->	Optional[Any]:
								__UpperCAmelCase							=			'''
        examples/by_feature/cross_validation.py
        --num_folds 2
        '''.split()
								with mock.patch.dict(os.environ     ,		{'''TESTING_MOCKED_DATALOADERS''': '''0'''}	):
									__UpperCAmelCase							=			run_command(self._launch_args + testargs     ,		return_stdout=__a	)
									__UpperCAmelCase							=			re.findall('''({.+})'''     ,		__a	)
									__UpperCAmelCase							=			[r for r in results if '''accuracy''' in r][-1]
									__UpperCAmelCase							=			ast.literal_eval(__a	)
									self.assertGreaterEqual(results['''accuracy''']     ,		0.7_5	)





							def       snake_case__						(      self				:							Dict	)							->	int:
								__UpperCAmelCase							=			['''examples/by_feature/multi_process_metrics.py''']
								run_command(self._launch_args + testargs	)





							@require_trackers
							@mock.patch.dict(os.environ     ,		{'''WANDB_MODE''': '''offline'''}	)
							def       snake_case__						(      self				:							Optional[Any]	)							->	Optional[int]:
								with tempfile.TemporaryDirectory() as tmpdir:
									__UpperCAmelCase							=			f"""
            examples/by_feature/tracking.py
            --with_tracking
            --project_dir {tmpdir}
            """.split()
									run_command(self._launch_args + testargs	)
									self.assertTrue(os.path.exists(os.path.join(__a     ,		'''tracking'''	)	)	)





							def       snake_case__						(      self				:							Optional[int]	)							->	List[Any]:
								__UpperCAmelCase							=			['''examples/by_feature/gradient_accumulation.py''']
								run_command(self._launch_args + testargs	)






							def       snake_case__						(      self				:							Tuple	)							->	Optional[Any]:
								__UpperCAmelCase							=			['''examples/by_feature/local_sgd.py''']
								run_command(self._launch_args + testargs	)
