





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__(__snake_case							)       ->   Optional[Any]:







	'''simple docstring'''







	lowerCamelCase__              =	[2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 18, 2]
	lowerCamelCase__              =	True if '''large''' in model_name or '''huge''' in model_name else False
	lowerCamelCase__              =	True if '''large''' in model_name or '''huge''' in model_name else False
	lowerCamelCase__              =	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:
			lowerCamelCase__              =	[3, 3, 3, 3]
			lowerCamelCase__              =	[5, 5, 5, 5]
		elif "fl4" in model_name:
			lowerCamelCase__              =	[4, 4, 4, 4]
			lowerCamelCase__              =	[3, 3, 3, 3]

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

	if "tiny" in model_name:
		lowerCamelCase__              =	96
	elif "small" in model_name:
		lowerCamelCase__              =	96
	elif "base" in model_name:
		lowerCamelCase__              =	128
	elif "large" in model_name:
		lowerCamelCase__              =	192
	elif "xlarge" in model_name:
		lowerCamelCase__              =	256
	elif "huge" in model_name:
		lowerCamelCase__              =	352

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

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

	lowerCamelCase__              =	FocalNetConfig(
	    embed_dim=__snake_case ,depths=__snake_case ,focal_levels=__snake_case ,focal_windows=__snake_case ,use_conv_embed=__snake_case ,idalabel=__snake_case ,labelaid=__snake_case ,use_post_layernorm=__snake_case ,use_layerscale=__snake_case ,)

	return config





def lowerCAmelCase__(__snake_case							)       ->   Optional[Any]:







	'''simple docstring'''







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

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

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

	return name





def lowerCAmelCase__(__snake_case ,__snake_case ,__snake_case=False							)       ->   Optional[int]:







	'''simple docstring'''







	lowerCamelCase__              =	{
	    '''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

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

	# rename keys
	for key in state_dict.copy().keys():
		lowerCamelCase__              =	state_dict.pop(__snake_case							)
		lowerCamelCase__              =	val

	lowerCamelCase__              =	get_focalnet_config(__snake_case							)
	lowerCamelCase__              =	FocalNetForImageClassification(__snake_case							)
	model.eval()

	# load state dict
	model.load_state_dict(__snake_case							)

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

	lowerCamelCase__              =	BitImageProcessor(
	    do_resize=__snake_case ,size={'''shortest_edge''': 256} ,resample=PILImageResampling.BILINEAR ,do_center_crop=__snake_case ,crop_size=224 ,do_normalize=__snake_case ,image_mean=__snake_case ,image_std=__snake_case ,)
	lowerCamelCase__              =	Image.open(requests.get(__snake_case ,stream=__snake_case							).raw							)
	lowerCamelCase__              =	processor(images=__snake_case ,return_tensors='''pt'''							)

	lowerCamelCase__              =	transforms.Compose(
	    [
	        transforms.Resize(256							),
	        transforms.CenterCrop(224							),
	        transforms.ToTensor(),
	        transforms.Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6] ,std=[0.2_2_9, 0.2_2_4, 0.2_2_5]							),
	    ]							)

	lowerCamelCase__              =	image_transforms(__snake_case							).unsqueeze(0							)

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

	lowerCamelCase__              =	model(**__snake_case							)

	lowerCamelCase__              =	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":
		lowerCamelCase__              =	torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]							)
	elif model_name == "focalnet-tiny-lrf":
		lowerCamelCase__              =	torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]							)
	elif model_name == "focalnet-small":
		lowerCamelCase__              =	torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]							)
	elif model_name == "focalnet-small-lrf":
		lowerCamelCase__              =	torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]							)
	elif model_name == "focalnet-base":
		lowerCamelCase__              =	torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]							)
	elif model_name == "focalnet-base-lrf":
		lowerCamelCase__              =	torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]							)
	assert torch.allclose(outputs.logits[0, :3] ,__snake_case ,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(__snake_case							)
		processor.save_pretrained(__snake_case							)

	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__":
			_a						=   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.",
			)

			_a						=   parser.parse_args()
			convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin






@dataclass
class       __A					(					lowerCAmelCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		42






class       __A					(					lowerCAmelCase    ,       lowerCAmelCase				):

			'''simple docstring'''







			@register_to_config
			def __init__(					self ,						__lowerCAmelCase = 1_6 ,						__lowerCAmelCase = 8_8 ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						__lowerCAmelCase = 1 ,						__lowerCAmelCase = 0.0 ,						__lowerCAmelCase = 3_2 ,						__lowerCAmelCase = None ,						__lowerCAmelCase = False ,						__lowerCAmelCase = None ,						__lowerCAmelCase = "geglu" ,						__lowerCAmelCase = True ,						__lowerCAmelCase = True ,						):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	num_attention_heads
				lowerCamelCase__              =	attention_head_dim
				lowerCamelCase__              =	num_attention_heads * attention_head_dim

				lowerCamelCase__              =	in_channels

				lowerCamelCase__              =	torch.nn.GroupNorm(num_groups=__lowerCAmelCase ,						num_channels=__lowerCAmelCase ,						eps=1E-6 ,						affine=__lowerCAmelCase							)
				lowerCamelCase__              =	nn.Linear(__lowerCAmelCase ,						__lowerCAmelCase							)

				# 3. Define transformers blocks
				lowerCamelCase__              =	nn.ModuleList(
				    [
				        BasicTransformerBlock(
				            __lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						dropout=__lowerCAmelCase ,						cross_attention_dim=__lowerCAmelCase ,						activation_fn=__lowerCAmelCase ,						attention_bias=__lowerCAmelCase ,						double_self_attention=__lowerCAmelCase ,						norm_elementwise_affine=__lowerCAmelCase ,						)
				        for d in range(__lowerCAmelCase							)
				    ]							)

				lowerCamelCase__              =	nn.Linear(__lowerCAmelCase ,						__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase=None ,						__lowerCAmelCase=None ,						__lowerCAmelCase=None ,						__lowerCAmelCase=1 ,						__lowerCAmelCase=None ,						__lowerCAmelCase = True ,						):



				'''simple docstring'''







				lowerCamelCase__       ,			lowerCamelCase__       ,			lowerCamelCase__       ,			lowerCamelCase__              =	hidden_states.shape
				lowerCamelCase__              =	batch_frames // num_frames

				lowerCamelCase__              =	hidden_states

				lowerCamelCase__              =	hidden_states[None, :].reshape(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							)
				lowerCamelCase__              =	hidden_states.permute(0 ,						2 ,						1 ,						3 ,						4							)

				lowerCamelCase__              =	self.norm(__lowerCAmelCase							)
				lowerCamelCase__              =	hidden_states.permute(0 ,						3 ,						4 ,						2 ,						1							).reshape(batch_size * height * width ,						__lowerCAmelCase ,						__lowerCAmelCase							)

				lowerCamelCase__              =	self.proj_in(__lowerCAmelCase							)

				# 2. Blocks
				for block in self.transformer_blocks:
					lowerCamelCase__              =	block(
					    __lowerCAmelCase ,						encoder_hidden_states=__lowerCAmelCase ,						timestep=__lowerCAmelCase ,						cross_attention_kwargs=__lowerCAmelCase ,						class_labels=__lowerCAmelCase ,						)

				# 3. Output
				lowerCamelCase__              =	self.proj_out(__lowerCAmelCase							)
				lowerCamelCase__              =	(
				    hidden_states[None, None, :]
				    .reshape(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							)
				    .permute(0 ,						3 ,						4 ,						1 ,						2							)
				    .contiguous()
				)
				lowerCamelCase__              =	hidden_states.reshape(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							)

				lowerCamelCase__              =	hidden_states + residual

				if not return_dict:
					return (output,)

				return TransformerTemporalModelOutput(sample=__lowerCAmelCase							)



