





import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

from transformers import is_torch_available
from transformers.testing_utils import is_flaky, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


SCREAMING_SNAKE_CASE_		     = 1e-4

if is_torch_available():
      import torch

      from transformers import AutoformerConfig, AutoformerForPrediction, AutoformerModel
      from transformers.models.autoformer.modeling_autoformer import AutoformerDecoder, AutoformerEncoder

@require_torch
class 	UpperCamelCase__				:



     '''simple docstring'''

     def __init__(							self		:					Tuple		,lowerCamelCase__		:					Dict		,lowerCamelCase__		:					Any=16		,lowerCamelCase__		:					str=13		,lowerCamelCase__		:					Union[str, Any]=7		,lowerCamelCase__		:					List[str]=14		,lowerCamelCase__		:					Tuple=10		,lowerCamelCase__		:					int=19		,lowerCamelCase__		:					Optional[int]=5		,lowerCamelCase__		:					str=4		,lowerCamelCase__		:					Union[str, Any]=True		,lowerCamelCase__		:					str=16		,lowerCamelCase__		:					int=2		,lowerCamelCase__		:					List[str]=4		,lowerCamelCase__		:					Union[str, Any]=4		,lowerCamelCase__		:					Optional[int]="gelu"		,lowerCamelCase__		:					Tuple=0.1		,lowerCamelCase__		:					Tuple=0.1		,lowerCamelCase__		:					int=[1, 2, 3, 4, 5]		,lowerCamelCase__		:					Any=25		,lowerCamelCase__		:					List[Any]=5		,)  ->	Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		d_model
           SCREAMING_SNAKE_CASE	=		parent
           SCREAMING_SNAKE_CASE	=		batch_size
           SCREAMING_SNAKE_CASE	=		prediction_length
           SCREAMING_SNAKE_CASE	=		context_length
           SCREAMING_SNAKE_CASE	=		cardinality
           SCREAMING_SNAKE_CASE	=		num_time_features
           SCREAMING_SNAKE_CASE	=		lags_sequence
           SCREAMING_SNAKE_CASE	=		embedding_dimension
           SCREAMING_SNAKE_CASE	=		is_training
           SCREAMING_SNAKE_CASE	=		hidden_size
           SCREAMING_SNAKE_CASE	=		num_hidden_layers
           SCREAMING_SNAKE_CASE	=		num_attention_heads
           SCREAMING_SNAKE_CASE	=		intermediate_size
           SCREAMING_SNAKE_CASE	=		hidden_act
           SCREAMING_SNAKE_CASE	=		hidden_dropout_prob
           SCREAMING_SNAKE_CASE	=		attention_probs_dropout_prob

           SCREAMING_SNAKE_CASE	=		context_length
           SCREAMING_SNAKE_CASE	=		prediction_length + label_length
           SCREAMING_SNAKE_CASE	=		label_length

           SCREAMING_SNAKE_CASE	=		moving_average
           SCREAMING_SNAKE_CASE	=		autocorrelation_factor






     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[int] )  ->	str:




           '''simple docstring'''


           return AutoformerConfig(
               d_model=self.d_model		,encoder_layers=self.num_hidden_layers		,decoder_layers=self.num_hidden_layers		,encoder_attention_heads=self.num_attention_heads		,decoder_attention_heads=self.num_attention_heads		,encoder_ffn_dim=self.intermediate_size		,decoder_ffn_dim=self.intermediate_size		,dropout=self.hidden_dropout_prob		,attention_dropout=self.attention_probs_dropout_prob		,prediction_length=self.prediction_length		,context_length=self.context_length		,label_length=self.label_length		,lags_sequence=self.lags_sequence		,num_time_features=self.num_time_features		,num_static_categorical_features=1		,cardinality=[self.cardinality]		,embedding_dimension=[self.embedding_dimension]		,moving_average=self.moving_average		,)






     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple		,lowerCamelCase__		:					int )  ->	Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		config.context_length + max(config.lags_sequence )

           SCREAMING_SNAKE_CASE	=		ids_tensor([self.batch_size, 1]		,config.cardinality[0] )
           SCREAMING_SNAKE_CASE	=		floats_tensor([self.batch_size, _past_length, config.num_time_features] )
           SCREAMING_SNAKE_CASE	=		floats_tensor([self.batch_size, _past_length] )
           SCREAMING_SNAKE_CASE	=		floats_tensor([self.batch_size, _past_length] ) > 0.5

           # decoder inputs
           SCREAMING_SNAKE_CASE	=		floats_tensor([self.batch_size, config.prediction_length, config.num_time_features] )
           SCREAMING_SNAKE_CASE	=		floats_tensor([self.batch_size, config.prediction_length] )

           SCREAMING_SNAKE_CASE	=		{
               """past_values""": past_values,
               """static_categorical_features""": static_categorical_features,
               """past_time_features""": past_time_features,
               """past_observed_mask""": past_observed_mask,
               """future_time_features""": future_time_features,
               """future_values""": future_values,
           }
           return inputs_dict






     def        SCREAMING_SNAKE_CASE__						(							self		:					Dict )  ->	Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		self.get_config()
           SCREAMING_SNAKE_CASE	=		self.prepare_autoformer_inputs_dict(lowerCamelCase__ )
           return config, inputs_dict






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[Any] )  ->	Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self.prepare_config_and_inputs()
           return config, inputs_dict






     def        SCREAMING_SNAKE_CASE__						(							self		:					int		,lowerCamelCase__		:					str		,lowerCamelCase__		:					List[Any] )  ->	Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		AutoformerModel(config=lowerCamelCase__ ).to(lowerCamelCase__ ).eval()
           SCREAMING_SNAKE_CASE	=		model(**lowerCamelCase__ )

           SCREAMING_SNAKE_CASE	=		outputs.encoder_last_hidden_state
           SCREAMING_SNAKE_CASE	=		outputs.last_hidden_state

           with tempfile.TemporaryDirectory() as tmpdirname:
                 SCREAMING_SNAKE_CASE	=		model.get_encoder()
                 encoder.save_pretrained(lowerCamelCase__ )
                 SCREAMING_SNAKE_CASE	=		AutoformerEncoder.from_pretrained(lowerCamelCase__ ).to(lowerCamelCase__ )

           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		model.create_network_inputs(**lowerCamelCase__ )
           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		model.decomposition_layer(transformer_inputs[:, : config.context_length, ...] )

           SCREAMING_SNAKE_CASE	=		torch.cat(
               (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...])		,dim=-1		,)
           SCREAMING_SNAKE_CASE	=		encoder(inputs_embeds=lowerCamelCase__ )[0]
           self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1e-3 )

           SCREAMING_SNAKE_CASE	=		(
               torch.mean(transformer_inputs[:, : config.context_length, ...]		,dim=1 )
               .unsqueeze(1 )
               .repeat(1		,config.prediction_length		,1 )
           )
           SCREAMING_SNAKE_CASE	=		torch.zeros(
               [transformer_inputs.shape[0], config.prediction_length, transformer_inputs.shape[2]]		,device=enc_input.device		,)

           SCREAMING_SNAKE_CASE	=		torch.cat(
               (
                   torch.cat((seasonal_input[:, -config.label_length :, ...], zeros)		,dim=1 ),
                   feature[:, config.context_length - config.label_length :, ...],
               )		,dim=-1		,)
           SCREAMING_SNAKE_CASE	=		torch.cat(
               (
                   torch.cat((trend_input[:, -config.label_length :, ...], mean)		,dim=1 ),
                   feature[:, config.context_length - config.label_length :, ...],
               )		,dim=-1		,)

           with tempfile.TemporaryDirectory() as tmpdirname:
                 SCREAMING_SNAKE_CASE	=		model.get_decoder()
                 decoder.save_pretrained(lowerCamelCase__ )
                 SCREAMING_SNAKE_CASE	=		AutoformerDecoder.from_pretrained(lowerCamelCase__ ).to(lowerCamelCase__ )

           SCREAMING_SNAKE_CASE	=		decoder(
               trend=lowerCamelCase__		,inputs_embeds=lowerCamelCase__		,encoder_hidden_states=lowerCamelCase__		,)[0]

           self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1e-3 )

@require_torch
class 	UpperCamelCase__				(   lowerCAmelCase_							,   lowerCAmelCase_							,   unittest.TestCase     ):



     '''simple docstring'''

     __snake_case					:			str				=    (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
     __snake_case					:			Union[str, Any]				=    (AutoformerForPrediction,) if is_torch_available() else ()
     __snake_case					:			str				=    {"feature-extraction": AutoformerModel} if is_torch_available() else {}
     __snake_case					:			Optional[Any]				=    False
     __snake_case					:			List[Any]				=    False
     __snake_case					:			Tuple				=    False
     __snake_case					:			Union[str, Any]				=    False
     __snake_case					:			Union[str, Any]				=    False
     __snake_case					:			Dict				=    False






     def        SCREAMING_SNAKE_CASE__						(							self		:					Union[str, Any] )  ->	Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		AutoformerModelTester(self )
           SCREAMING_SNAKE_CASE	=		ConfigTester(self		,config_class=lowerCamelCase__		,has_text_modality=lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					Dict )  ->	int:




           '''simple docstring'''


           self.config_tester.run_common_tests()






     def        SCREAMING_SNAKE_CASE__						(							self		:					str )  ->	Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
                 SCREAMING_SNAKE_CASE	=		model_class(lowerCamelCase__ )

                 with tempfile.TemporaryDirectory() as tmpdirname:
                       model.save_pretrained(lowerCamelCase__ )
                       SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		model_class.from_pretrained(lowerCamelCase__		,output_loading_info=lowerCamelCase__ )
                 self.assertEqual(info["""missing_keys"""]		,[] )






     def        SCREAMING_SNAKE_CASE__						(							self		:					Union[str, Any] )  ->	List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		self.model_tester.prepare_config_and_inputs_for_common()
           self.model_tester.check_encoder_decoder_model_standalone(*lowerCamelCase__ )






     @unittest.skip(reason="""Model has no tokens embeddings""" )
     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple )  ->	Optional[int]:




           '''simple docstring'''


           pass






     def        SCREAMING_SNAKE_CASE__						(							self		:					int )  ->	Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		inspect.signature(getattr(lowerCamelCase__		,"""forward""" ) )
           # The main input is the name of the argument after `self`
           SCREAMING_SNAKE_CASE	=		list(model_signature.parameters.keys() )[1]
           self.assertEqual(AutoformerModel.main_input_name		,lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					int )  ->	Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 SCREAMING_SNAKE_CASE	=		model_class(lowerCamelCase__ )
                 SCREAMING_SNAKE_CASE	=		inspect.signature(model.forward )
                 # signature.parameters is an OrderedDict => so arg_names order is deterministic
                 SCREAMING_SNAKE_CASE	=		[*signature.parameters.keys()]

                 SCREAMING_SNAKE_CASE	=		[
                     """past_values""",
                     """past_time_features""",
                     """past_observed_mask""",
                     """static_categorical_features""",
                     """static_real_features""",
                     """future_values""",
                     """future_time_features""",
                 ]

                 if model.__class__.__name__ in ["AutoformerForPrediction"]:
                       expected_arg_names.append("""future_observed_mask""" )

                 expected_arg_names.extend(
                     [
                         """decoder_attention_mask""",
                         """head_mask""",
                         """decoder_head_mask""",
                         """cross_attn_head_mask""",
                         """encoder_outputs""",
                         """past_key_values""",
                         """output_hidden_states""",
                         """output_attentions""",
                         """use_cache""",
                         """return_dict""",
                     ] )

                 self.assertListEqual(arg_names[: len(lowerCamelCase__ )]		,lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[int] )  ->	Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self.model_tester.prepare_config_and_inputs_for_common()
           SCREAMING_SNAKE_CASE	=		True

           SCREAMING_SNAKE_CASE	=		getattr(self.model_tester		,"""seq_length"""		,lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		getattr(self.model_tester		,"""decoder_seq_length"""		,lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		getattr(self.model_tester		,"""encoder_seq_length"""		,lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		getattr(self.model_tester		,"""d_model"""		,lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		getattr(self.model_tester		,"""num_attention_heads"""		,lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		d_model // num_attention_heads

           for model_class in self.all_model_classes:
                 SCREAMING_SNAKE_CASE	=		True
                 SCREAMING_SNAKE_CASE	=		False
                 SCREAMING_SNAKE_CASE	=		True
                 SCREAMING_SNAKE_CASE	=		model_class(lowerCamelCase__ )
                 model.to(lowerCamelCase__ )
                 model.eval()
                 with torch.no_grad():
                       SCREAMING_SNAKE_CASE	=		model(**self._prepare_for_class(lowerCamelCase__		,lowerCamelCase__ ) )
                 SCREAMING_SNAKE_CASE	=		outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
                 self.assertEqual(len(lowerCamelCase__ )		,self.model_tester.num_hidden_layers )

                 # check that output_attentions also work using config
                 del inputs_dict["output_attentions"]
                 SCREAMING_SNAKE_CASE	=		True
                 SCREAMING_SNAKE_CASE	=		model_class(lowerCamelCase__ )
                 model.to(lowerCamelCase__ )
                 model.eval()
                 with torch.no_grad():
                       SCREAMING_SNAKE_CASE	=		model(**self._prepare_for_class(lowerCamelCase__		,lowerCamelCase__ ) )
                 SCREAMING_SNAKE_CASE	=		outputs.encoder_attentions
                 self.assertEqual(len(lowerCamelCase__ )		,self.model_tester.num_hidden_layers )

                 self.assertListEqual(
                     list(attentions[0].shape[-3:] )		,[self.model_tester.num_attention_heads, encoder_seq_length, dim]		,)
                 SCREAMING_SNAKE_CASE	=		len(lowerCamelCase__ )

                 SCREAMING_SNAKE_CASE	=		7

                 if "last_hidden_state" in outputs:
                       correct_outlen += 1

                 if "trend" in outputs:
                       correct_outlen += 1

                 if "past_key_values" in outputs:
                       correct_outlen += 1  # past_key_values have been returned

                 if "loss" in outputs:
                       correct_outlen += 1

                 if "params" in outputs:
                       correct_outlen += 1

                 self.assertEqual(lowerCamelCase__		,lowerCamelCase__ )

                 # decoder attentions
                 SCREAMING_SNAKE_CASE	=		outputs.decoder_attentions
                 self.assertIsInstance(lowerCamelCase__		,(list, tuple) )
                 self.assertEqual(len(lowerCamelCase__ )		,self.model_tester.num_hidden_layers )
                 self.assertListEqual(
                     list(decoder_attentions[0].shape[-3:] )		,[self.model_tester.num_attention_heads, decoder_seq_length, dim]		,)

                 # cross attentions
                 SCREAMING_SNAKE_CASE	=		outputs.cross_attentions
                 self.assertIsInstance(lowerCamelCase__		,(list, tuple) )
                 self.assertEqual(len(lowerCamelCase__ )		,self.model_tester.num_hidden_layers )
                 self.assertListEqual(
                     list(cross_attentions[0].shape[-3:] )		,[self.model_tester.num_attention_heads, decoder_seq_length, dim]		,)

           # Check attention is always last and order is fine
           SCREAMING_SNAKE_CASE	=		True
           SCREAMING_SNAKE_CASE	=		True
           SCREAMING_SNAKE_CASE	=		model_class(lowerCamelCase__ )
           model.to(lowerCamelCase__ )
           model.eval()
           with torch.no_grad():
                 SCREAMING_SNAKE_CASE	=		model(**self._prepare_for_class(lowerCamelCase__		,lowerCamelCase__ ) )

           self.assertEqual(out_len + 2		,len(lowerCamelCase__ ) )

           SCREAMING_SNAKE_CASE	=		outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

           self.assertEqual(len(lowerCamelCase__ )		,self.model_tester.num_hidden_layers )
           self.assertListEqual(
               list(self_attentions[0].shape[-3:] )		,[self.model_tester.num_attention_heads, encoder_seq_length, dim]		,)






     @is_flaky()
     def        SCREAMING_SNAKE_CASE__						(							self		:					Dict )  ->	List[str]:




           '''simple docstring'''


           super().test_retain_grad_hidden_states_attentions()



def      __lowercase     (   _SCREAMING_SNAKE_CASE="train-batch.pt"			)				->     Any:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE	=		hf_hub_download(repo_id="""hf-internal-testing/tourism-monthly-batch"""				,		filename=_SCREAMING_SNAKE_CASE				,		repo_type="""dataset"""			)
      SCREAMING_SNAKE_CASE	=		torch.load(_SCREAMING_SNAKE_CASE				,		map_location=_SCREAMING_SNAKE_CASE			)
      return batch



@require_torch
@slow
class 	UpperCamelCase__				(   unittest.TestCase     ):



     '''simple docstring'''

     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple )  ->	str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		AutoformerModel.from_pretrained("""huggingface/autoformer-tourism-monthly""" ).to(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		prepare_batch()

           with torch.no_grad():
                 SCREAMING_SNAKE_CASE	=		model(
                     past_values=batch["""past_values"""]		,past_time_features=batch["""past_time_features"""]		,past_observed_mask=batch["""past_observed_mask"""]		,static_categorical_features=batch["""static_categorical_features"""]		,future_values=batch["""future_values"""]		,future_time_features=batch["""future_time_features"""]		,)[0]

           SCREAMING_SNAKE_CASE	=		torch.Size(
               (64, model.config.prediction_length + model.config.label_length, model.config.feature_size) )
           self.assertEqual(output.shape		,lowerCamelCase__ )

           SCREAMING_SNAKE_CASE	=		torch.tensor(
               [[0.3593, -1.3398, 0.6330], [0.2279, 1.5396, -0.1792], [0.0450, 1.3225, -0.2335]]		,device=lowerCamelCase__ )
           self.assertTrue(torch.allclose(output[0, :3, :3]		,lowerCamelCase__		,atol=lowerCamelCase__ ) )






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[str] )  ->	str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly""" ).to(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		prepare_batch("""val-batch.pt""" )
           with torch.no_grad():
                 SCREAMING_SNAKE_CASE	=		model(
                     past_values=batch["""past_values"""]		,past_time_features=batch["""past_time_features"""]		,past_observed_mask=batch["""past_observed_mask"""]		,static_categorical_features=batch["""static_categorical_features"""]		,).encoder_last_hidden_state
           SCREAMING_SNAKE_CASE	=		torch.Size((64, model.config.context_length, model.config.d_model) )
           self.assertEqual(output.shape		,lowerCamelCase__ )

           SCREAMING_SNAKE_CASE	=		torch.tensor(
               [[-0.0734, -0.9036, 0.8358], [4.7186, 2.4113, 1.9581], [1.7953, 2.3558, 1.2970]]		,device=lowerCamelCase__ )
           self.assertTrue(torch.allclose(output[0, :3, :3]		,lowerCamelCase__		,atol=lowerCamelCase__ ) )






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[str] )  ->	Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly""" ).to(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		prepare_batch("""val-batch.pt""" )
           with torch.no_grad():
                 SCREAMING_SNAKE_CASE	=		model.generate(
                     static_categorical_features=batch["""static_categorical_features"""]		,past_time_features=batch["""past_time_features"""]		,past_values=batch["""past_values"""]		,future_time_features=batch["""future_time_features"""]		,past_observed_mask=batch["""past_observed_mask"""]		,)
           SCREAMING_SNAKE_CASE	=		torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length) )
           self.assertEqual(outputs.sequences.shape		,lowerCamelCase__ )

           SCREAMING_SNAKE_CASE	=		torch.tensor([3130.6763, 4056.5293, 7053.0786]		,device=lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		outputs.sequences.mean(dim=1 )
           self.assertTrue(torch.allclose(mean_prediction[0, -3:]		,lowerCamelCase__		,rtol=1e-1 ) )









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      __lowercase     (   _SCREAMING_SNAKE_CASE			)				->     Any:


      '''simple docstring'''


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

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

      if "tiny" in model_name:
            SCREAMING_SNAKE_CASE	=		96
      elif "small" in model_name:
            SCREAMING_SNAKE_CASE	=		96
      elif "base" in model_name:
            SCREAMING_SNAKE_CASE	=		1_28
      elif "large" in model_name:
            SCREAMING_SNAKE_CASE	=		1_92
      elif "xlarge" in model_name:
            SCREAMING_SNAKE_CASE	=		2_56
      elif "huge" in model_name:
            SCREAMING_SNAKE_CASE	=		3_52

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

      SCREAMING_SNAKE_CASE	=		json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE				,		repo_type="""dataset"""			)				,		"""r"""			)			)
      SCREAMING_SNAKE_CASE	=		{int(_SCREAMING_SNAKE_CASE			): v for k, v in idalabel.items()}
      SCREAMING_SNAKE_CASE	=		{v: k for k, v in idalabel.items()}

      SCREAMING_SNAKE_CASE	=		FocalNetConfig(
          embed_dim=_SCREAMING_SNAKE_CASE				,		depths=_SCREAMING_SNAKE_CASE				,		focal_levels=_SCREAMING_SNAKE_CASE				,		focal_windows=_SCREAMING_SNAKE_CASE				,		use_conv_embed=_SCREAMING_SNAKE_CASE				,		idalabel=_SCREAMING_SNAKE_CASE				,		labelaid=_SCREAMING_SNAKE_CASE				,		use_post_layernorm=_SCREAMING_SNAKE_CASE				,		use_layerscale=_SCREAMING_SNAKE_CASE				,		)

      return config



def      __lowercase     (   _SCREAMING_SNAKE_CASE			)				->     Optional[int]:


      '''simple docstring'''


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

      if name == "norm.weight":
            SCREAMING_SNAKE_CASE	=		"""layernorm.weight"""
      if name == "norm.bias":
            SCREAMING_SNAKE_CASE	=		"""layernorm.bias"""

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

      return name



def      __lowercase     (   _SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE=False			)				->     Dict:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE	=		{
          """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

      SCREAMING_SNAKE_CASE	=		model_name_to_url[model_name]
      print("""Checkpoint URL: """				,		_SCREAMING_SNAKE_CASE			)
      SCREAMING_SNAKE_CASE	=		torch.hub.load_state_dict_from_url(_SCREAMING_SNAKE_CASE				,		map_location="""cpu"""			)["""model"""]

      # rename keys
      for key in state_dict.copy().keys():
            SCREAMING_SNAKE_CASE	=		state_dict.pop(_SCREAMING_SNAKE_CASE			)
            SCREAMING_SNAKE_CASE	=		val

      SCREAMING_SNAKE_CASE	=		get_focalnet_config(_SCREAMING_SNAKE_CASE			)
      SCREAMING_SNAKE_CASE	=		FocalNetForImageClassification(_SCREAMING_SNAKE_CASE			)
      model.eval()

      # load state dict
      model.load_state_dict(_SCREAMING_SNAKE_CASE			)

      # verify conversion
      SCREAMING_SNAKE_CASE	=		"""http://images.cocodataset.org/val2017/000000039769.jpg"""

      SCREAMING_SNAKE_CASE	=		BitImageProcessor(
          do_resize=_SCREAMING_SNAKE_CASE				,		size={"""shortest_edge""": 2_56}				,		resample=PILImageResampling.BILINEAR				,		do_center_crop=_SCREAMING_SNAKE_CASE				,		crop_size=2_24				,		do_normalize=_SCREAMING_SNAKE_CASE				,		image_mean=_SCREAMING_SNAKE_CASE				,		image_std=_SCREAMING_SNAKE_CASE				,		)
      SCREAMING_SNAKE_CASE	=		Image.open(requests.get(_SCREAMING_SNAKE_CASE				,		stream=_SCREAMING_SNAKE_CASE			).raw			)
      SCREAMING_SNAKE_CASE	=		processor(images=_SCREAMING_SNAKE_CASE				,		return_tensors="""pt"""			)

      SCREAMING_SNAKE_CASE	=		transforms.Compose(
          [
              transforms.Resize(2_56			),
              transforms.CenterCrop(2_24			),
              transforms.ToTensor(),
              transforms.Normalize(mean=[0.485, 0.456, 0.406]				,		std=[0.229, 0.224, 0.225]			),
          ]			)

      SCREAMING_SNAKE_CASE	=		image_transforms(_SCREAMING_SNAKE_CASE			).unsqueeze(0			)

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

      SCREAMING_SNAKE_CASE	=		model(**_SCREAMING_SNAKE_CASE			)

      SCREAMING_SNAKE_CASE	=		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":
            SCREAMING_SNAKE_CASE	=		torch.tensor([0.2_166, -0.4_368, 0.2_191]			)
      elif model_name == "focalnet-tiny-lrf":
            SCREAMING_SNAKE_CASE	=		torch.tensor([1.1_669, 0.0_125, -0.1_695]			)
      elif model_name == "focalnet-small":
            SCREAMING_SNAKE_CASE	=		torch.tensor([0.4_917, -0.0_430, 0.1_341]			)
      elif model_name == "focalnet-small-lrf":
            SCREAMING_SNAKE_CASE	=		torch.tensor([-0.2_588, -0.5_342, -0.2_331]			)
      elif model_name == "focalnet-base":
            SCREAMING_SNAKE_CASE	=		torch.tensor([-0.1_655, -0.4_090, -0.1_730]			)
      elif model_name == "focalnet-base-lrf":
            SCREAMING_SNAKE_CASE	=		torch.tensor([0.5_306, -0.0_483, -0.3_928]			)
      assert torch.allclose(outputs.logits[0, :3]				,		_SCREAMING_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(_SCREAMING_SNAKE_CASE			)
            processor.save_pretrained(_SCREAMING_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__":
      SCREAMING_SNAKE_CASE_		     = 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.""",
      )

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


