
'''simple docstring'''



import inspect
import unittest

from transformers import SwinaSRConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

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


if is_torch_available():
     import torch
     from torch import nn

     from transformers import SwinaSRForImageSuperResolution, SwinaSRModel
     from transformers.models.swinasr.modeling_swinasr import SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
     from PIL import Image

     from transformers import SwinaSRImageProcessor

class       _a :





       """simple docstring"""




       def __init__(							self				,							A__				,							A__=13				,							A__=32				,							A__=1				,							A__=3				,							A__=16				,							A__=[1, 2, 1]				,							A__=[2, 2, 4]				,							A__=2				,							A__=2.0				,							A__=True				,							A__=0.0				,							A__=0.0				,							A__=0.1				,							A__="gelu"				,							A__=False				,							A__=True				,							A__=0.02				,							A__=1E-5				,							A__=True				,							A__=None				,							A__=False				,							A__=2				,							)			->	Optional[Any]:
          _SCREAMING_SNAKE_CASE          = parent
          _SCREAMING_SNAKE_CASE          = batch_size
          _SCREAMING_SNAKE_CASE          = image_size
          _SCREAMING_SNAKE_CASE          = patch_size
          _SCREAMING_SNAKE_CASE          = num_channels
          _SCREAMING_SNAKE_CASE          = embed_dim
          _SCREAMING_SNAKE_CASE          = depths
          _SCREAMING_SNAKE_CASE          = num_heads
          _SCREAMING_SNAKE_CASE          = window_size
          _SCREAMING_SNAKE_CASE          = mlp_ratio
          _SCREAMING_SNAKE_CASE          = qkv_bias
          _SCREAMING_SNAKE_CASE          = hidden_dropout_prob
          _SCREAMING_SNAKE_CASE          = attention_probs_dropout_prob
          _SCREAMING_SNAKE_CASE          = drop_path_rate
          _SCREAMING_SNAKE_CASE          = hidden_act
          _SCREAMING_SNAKE_CASE          = use_absolute_embeddings
          _SCREAMING_SNAKE_CASE          = patch_norm
          _SCREAMING_SNAKE_CASE          = layer_norm_eps
          _SCREAMING_SNAKE_CASE          = initializer_range
          _SCREAMING_SNAKE_CASE          = is_training
          _SCREAMING_SNAKE_CASE          = scope
          _SCREAMING_SNAKE_CASE          = use_labels
          _SCREAMING_SNAKE_CASE          = upscale

          # here we set some attributes to make tests pass
          _SCREAMING_SNAKE_CASE          = len(A__					)
          _SCREAMING_SNAKE_CASE          = embed_dim
          _SCREAMING_SNAKE_CASE          = (image_size // patch_size) ** 2






       def 				UpperCamelCase (							self					)			->	Any:
          _SCREAMING_SNAKE_CASE          = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

          _SCREAMING_SNAKE_CASE          = None
          if self.use_labels:
             _SCREAMING_SNAKE_CASE          = ids_tensor([self.batch_size]				,							self.type_sequence_label_size					)

          _SCREAMING_SNAKE_CASE          = self.get_config()

          return config, pixel_values, labels






       def 				UpperCamelCase (							self					)			->	Optional[Any]:
          return SwinaSRConfig(
              image_size=self.image_size				,							patch_size=self.patch_size				,							num_channels=self.num_channels				,							embed_dim=self.embed_dim				,							depths=self.depths				,							num_heads=self.num_heads				,							window_size=self.window_size				,							mlp_ratio=self.mlp_ratio				,							qkv_bias=self.qkv_bias				,							hidden_dropout_prob=self.hidden_dropout_prob				,							attention_probs_dropout_prob=self.attention_probs_dropout_prob				,							drop_path_rate=self.drop_path_rate				,							hidden_act=self.hidden_act				,							use_absolute_embeddings=self.use_absolute_embeddings				,							path_norm=self.patch_norm				,							layer_norm_eps=self.layer_norm_eps				,							initializer_range=self.initializer_range				,							upscale=self.upscale				,							)






       def 				UpperCamelCase (							self				,							A__				,							A__				,							A__					)			->	List[str]:
          _SCREAMING_SNAKE_CASE          = SwinaSRModel(config=A__					)
          model.to(A__					)
          model.eval()
          _SCREAMING_SNAKE_CASE          = model(A__					)

          self.parent.assertEqual(
              result.last_hidden_state.shape				,							(self.batch_size, self.embed_dim, self.image_size, self.image_size)					)






       def 				UpperCamelCase (							self				,							A__				,							A__				,							A__					)			->	Dict:
          _SCREAMING_SNAKE_CASE          = SwinaSRForImageSuperResolution(A__					)
          model.to(A__					)
          model.eval()
          _SCREAMING_SNAKE_CASE          = model(A__					)

          _SCREAMING_SNAKE_CASE          = self.image_size * self.upscale
          self.parent.assertEqual(
              result.reconstruction.shape				,							(self.batch_size, self.num_channels, expected_image_size, expected_image_size)					)






       def 				UpperCamelCase (							self					)			->	Optional[Any]:
          _SCREAMING_SNAKE_CASE          = self.prepare_config_and_inputs()
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = config_and_inputs
          _SCREAMING_SNAKE_CASE          = {"""pixel_values""": pixel_values}
          return config, inputs_dict






@require_torch
class       _a (_lowerCamelCase     ,    _lowerCamelCase     ,    unittest.TestCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE					=     (SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
       SCREAMING_SNAKE_CASE					=     {'feature-extraction': SwinaSRModel} if is_torch_available() else {}

       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False






       def 				UpperCamelCase (							self					)			->	List[str]:
          _SCREAMING_SNAKE_CASE          = SwinaSRModelTester(self					)
          _SCREAMING_SNAKE_CASE          = ConfigTester(self				,							config_class=A__				,							embed_dim=37					)






       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.create_and_test_config_with_num_labels()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()






       def 				UpperCamelCase (							self					)			->	int:
          _SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*A__					)






       def 				UpperCamelCase (							self					)			->	List[Any]:
          _SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_super_resolution(*A__					)






       @unittest.skip(reason="""Got `CUDA error: misaligned address` with PyTorch 2.0.0."""					)
       def 				UpperCamelCase (							self					)			->	List[Any]:
          pass






       @unittest.skip(reason="""Swin2SR does not use inputs_embeds"""					)
       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          pass






       @unittest.skip(reason="""Swin2SR does not support training yet"""					)
       def 				UpperCamelCase (							self					)			->	Any:
          pass






       @unittest.skip(reason="""Swin2SR does not support training yet"""					)
       def 				UpperCamelCase (							self					)			->	Tuple:
          pass






       def 				UpperCamelCase (							self					)			->	Dict:
          _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(A__					)
             self.assertIsInstance(model.get_input_embeddings()				,							(nn.Module)					)
             _SCREAMING_SNAKE_CASE          = model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(A__				,							nn.Linear					)					)






       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          _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(A__					)
             _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          = ["""pixel_values"""]
             self.assertListEqual(arg_names[:1]				,							A__					)






       @slow
       def 				UpperCamelCase (							self					)			->	Optional[int]:
          for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _SCREAMING_SNAKE_CASE          = SwinaSRModel.from_pretrained(A__					)
             self.assertIsNotNone(A__					)






       def 				UpperCamelCase (							self					)			->	List[str]:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()

          _SCREAMING_SNAKE_CASE          = _config_zero_init(A__					)
          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE          = model_class(config=A__					)
             for name, param in model.named_parameters():
                if "logit_scale" in name:
                   continue
                if param.requires_grad:
                   self.assertIn(
                       ((param.data.mean() * 1E9).round() / 1E9).item()				,							[0.0, 1.0]				,							msg=F"Parameter {name} of model {model_class} seems not properly initialized"				,							)






       def 				UpperCamelCase (							self					)			->	Any:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()
          _SCREAMING_SNAKE_CASE          = True

          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(A__					)
             model.to(A__					)
             model.eval()
             with torch.no_grad():
                _SCREAMING_SNAKE_CASE          = model(**self._prepare_for_class(A__				,							A__					)					)
             _SCREAMING_SNAKE_CASE          = outputs.attentions
             _SCREAMING_SNAKE_CASE          = len(self.model_tester.depths					)
             self.assertEqual(len(A__					)				,							A__					)

             # check that output_attentions also work using config
             del inputs_dict["output_attentions"]
             _SCREAMING_SNAKE_CASE          = True
             _SCREAMING_SNAKE_CASE          = config.window_size**2
             _SCREAMING_SNAKE_CASE          = model_class(A__					)
             model.to(A__					)
             model.eval()
             with torch.no_grad():
                _SCREAMING_SNAKE_CASE          = model(**self._prepare_for_class(A__				,							A__					)					)
             _SCREAMING_SNAKE_CASE          = outputs.attentions
             self.assertEqual(len(A__					)				,							A__					)

             self.assertListEqual(
                 list(attentions[0].shape[-3:]					)				,							[self.model_tester.num_heads[0], window_size_squared, window_size_squared]				,							)
             _SCREAMING_SNAKE_CASE          = len(A__					)

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

             self.assertEqual(out_len + 1				,							len(A__					)					)

             _SCREAMING_SNAKE_CASE          = outputs.attentions

             self.assertEqual(len(A__					)				,							A__					)

             self.assertListEqual(
                 list(self_attentions[0].shape[-3:]					)				,							[self.model_tester.num_heads[0], window_size_squared, window_size_squared]				,							)






@require_vision
@require_torch
@slow
class       _a (unittest.TestCase):





       """simple docstring"""




       def 				UpperCamelCase (							self					)			->	Dict:
          _SCREAMING_SNAKE_CASE          = SwinaSRImageProcessor()
          _SCREAMING_SNAKE_CASE          = SwinaSRForImageSuperResolution.from_pretrained("""caidas/swin2SR-classical-sr-x2-64"""					).to(A__					)

          _SCREAMING_SNAKE_CASE          = Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
          _SCREAMING_SNAKE_CASE          = processor(images=A__				,							return_tensors="""pt"""					).to(A__					)

          # forward pass
          with torch.no_grad():
             _SCREAMING_SNAKE_CASE          = model(**A__					)

          # verify the logits
          _SCREAMING_SNAKE_CASE          = torch.Size([1, 3, 9_76, 12_96]					)
          self.assertEqual(outputs.reconstruction.shape				,							A__					)
          _SCREAMING_SNAKE_CASE          = torch.tensor(
              [[0.5458, 0.5546, 0.5638], [0.5526, 0.5565, 0.5651], [0.5396, 0.5426, 0.5621]]					).to(A__					)
          self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]				,							A__				,							atol=1E-4					)					)







'''simple docstring'''


import argparse
import json
from pathlib import Path

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

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
snake_case__     :			Optional[Any]      			=  logging.get_logger(__name__)


def 					_lowerCamelCase    (      lowerCamelCase_					:							List[str]   ):



     """simple docstring"""





     UpperCAmelCase_							:   Dict   =      ConvNextConfig()

     if "tiny" in checkpoint_url:
          UpperCAmelCase_							:   Any   =      [3, 3, 9, 3]
          UpperCAmelCase_							:   Any   =      [96, 192, 384, 768]
     if "small" in checkpoint_url:
          UpperCAmelCase_							:   int   =      [3, 3, 27, 3]
          UpperCAmelCase_							:   Optional[int]   =      [96, 192, 384, 768]
     if "base" in checkpoint_url:
          UpperCAmelCase_							:   str   =      [3, 3, 27, 3]
          UpperCAmelCase_							:   Dict   =      [128, 256, 512, 1024]
     if "large" in checkpoint_url:
          UpperCAmelCase_							:   Optional[Any]   =      [3, 3, 27, 3]
          UpperCAmelCase_							:   List[Any]   =      [192, 384, 768, 1536]
     if "xlarge" in checkpoint_url:
          UpperCAmelCase_							:   str   =      [3, 3, 27, 3]
          UpperCAmelCase_							:   Any   =      [256, 512, 1024, 2048]

     if "1k" in checkpoint_url:
          UpperCAmelCase_							:   int   =      1000
          UpperCAmelCase_							:   Any   =      'imagenet-1k-id2label.json'
          UpperCAmelCase_							:   str   =      (1, 1000)
     else:
          UpperCAmelCase_							:   Dict   =      21841
          UpperCAmelCase_							:   Optional[Any]   =      'imagenet-22k-id2label.json'
          UpperCAmelCase_							:   Optional[Any]   =      (1, 21841)

     UpperCAmelCase_							:   int   =      'huggingface/label-files'
     UpperCAmelCase_							:   Any   =      num_labels
     UpperCAmelCase_							:   Dict   =      json.load(open(hf_hub_download(lowerCamelCase_	,     lowerCamelCase_	,     repo_type='dataset'   )	,     'r'   )   )
     UpperCAmelCase_							:   List[Any]   =      {int(lowerCamelCase_   ): v for k, v in idalabel.items()}
     if "1k" not in checkpoint_url:
          # this dataset contains 21843 labels but the model only has 21841
          # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
          del idalabel[9205]
          del idalabel[15027]
     UpperCAmelCase_							:   Optional[Any]   =      idalabel
     UpperCAmelCase_							:   Dict   =      {v: k for k, v in idalabel.items()}
     UpperCAmelCase_							:   Tuple   =      hidden_sizes
     UpperCAmelCase_							:   List[str]   =      depths

     return config, expected_shape


def 					_lowerCamelCase    (      lowerCamelCase_					:							int   ):



     """simple docstring"""





     if "downsample_layers.0.0" in name:
          UpperCAmelCase_							:   List[Any]   =      name.replace('downsample_layers.0.0'	,     'embeddings.patch_embeddings'   )
     if "downsample_layers.0.1" in name:
          UpperCAmelCase_							:   List[Any]   =      name.replace('downsample_layers.0.1'	,     'embeddings.norm'   )  # we rename to layernorm later on
     if "downsample_layers.1.0" in name:
          UpperCAmelCase_							:   int   =      name.replace('downsample_layers.1.0'	,     'stages.1.downsampling_layer.0'   )
     if "downsample_layers.1.1" in name:
          UpperCAmelCase_							:   List[str]   =      name.replace('downsample_layers.1.1'	,     'stages.1.downsampling_layer.1'   )
     if "downsample_layers.2.0" in name:
          UpperCAmelCase_							:   str   =      name.replace('downsample_layers.2.0'	,     'stages.2.downsampling_layer.0'   )
     if "downsample_layers.2.1" in name:
          UpperCAmelCase_							:   List[str]   =      name.replace('downsample_layers.2.1'	,     'stages.2.downsampling_layer.1'   )
     if "downsample_layers.3.0" in name:
          UpperCAmelCase_							:   Optional[int]   =      name.replace('downsample_layers.3.0'	,     'stages.3.downsampling_layer.0'   )
     if "downsample_layers.3.1" in name:
          UpperCAmelCase_							:   Dict   =      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.
          UpperCAmelCase_							:   str   =      name[: len('stages.0'   )] + '.layers' + name[len('stages.0'   ) :]
     if "stages" in name:
          UpperCAmelCase_							:   int   =      name.replace('stages'	,     'encoder.stages'   )
     if "norm" in name:
          UpperCAmelCase_							:   List[str]   =      name.replace('norm'	,     'layernorm'   )
     if "gamma" in name:
          UpperCAmelCase_							:   Any   =      name.replace('gamma'	,     'layer_scale_parameter'   )
     if "head" in name:
          UpperCAmelCase_							:   List[str]   =      name.replace('head'	,     'classifier'   )

     return name


def 					_lowerCamelCase    (      ):



     """simple docstring"""





     UpperCAmelCase_							:   List[str]   =      'http://images.cocodataset.org/val2017/000000039769.jpg'
     UpperCAmelCase_							:   Optional[int]   =      Image.open(requests.get(lowerCamelCase_	,     stream=lowerCamelCase_   ).raw   )
     return im


@torch.no_grad()
def 					_lowerCamelCase    (      lowerCamelCase_					:							int	,     lowerCamelCase_					:							Optional[int]   ):



     """simple docstring"""





     UpperCAmelCase_   ,			UpperCAmelCase_							:   Optional[Any]   =      get_convnext_config(lowerCamelCase_   )
     # load original state_dict from URL
     UpperCAmelCase_							:   Dict   =      torch.hub.load_state_dict_from_url(lowerCamelCase_   )['model']
     # rename keys
     for key in state_dict.copy().keys():
          UpperCAmelCase_							:   Optional[int]   =      state_dict.pop(lowerCamelCase_   )
          UpperCAmelCase_							:   Tuple   =      val
     # add prefix to all keys expect classifier head
     for key in state_dict.copy().keys():
          UpperCAmelCase_							:   int   =      state_dict.pop(lowerCamelCase_   )
          if not key.startswith('classifier'   ):
               UpperCAmelCase_							:   int   =      'convnext.' + key
          UpperCAmelCase_							:   Tuple   =      val

     # load HuggingFace model
     UpperCAmelCase_							:   Tuple   =      ConvNextForImageClassification(lowerCamelCase_   )
     model.load_state_dict(lowerCamelCase_   )
     model.eval()

     # Check outputs on an image, prepared by ConvNextImageProcessor
     UpperCAmelCase_							:   Tuple   =      224 if '224' in checkpoint_url else 384
     UpperCAmelCase_							:   Optional[Any]   =      ConvNextImageProcessor(size=lowerCamelCase_   )
     UpperCAmelCase_							:   Tuple   =      image_processor(images=prepare_img()	,     return_tensors='pt'   ).pixel_values

     UpperCAmelCase_							:   Dict   =      model(lowerCamelCase_   ).logits

     # note: the logits below were obtained without center cropping
     if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
          UpperCAmelCase_							:   Tuple   =      torch.tensor([-0.1210, -0.6605, 0.1918]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
          UpperCAmelCase_							:   str   =      torch.tensor([-0.4473, -0.1847, -0.6365]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
          UpperCAmelCase_							:   Tuple   =      torch.tensor([0.4525, 0.7539, 0.0308]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
          UpperCAmelCase_							:   List[Any]   =      torch.tensor([0.3561, 0.6350, -0.0384]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
          UpperCAmelCase_							:   List[str]   =      torch.tensor([0.4174, -0.0989, 0.1489]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
          UpperCAmelCase_							:   Any   =      torch.tensor([0.2513, -0.1349, -0.1613]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
          UpperCAmelCase_							:   Dict   =      torch.tensor([1.2980, 0.3631, -0.1198]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
          UpperCAmelCase_							:   Union[str, Any]   =      torch.tensor([1.2963, 0.1227, 0.1723]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
          UpperCAmelCase_							:   Optional[Any]   =      torch.tensor([1.7956, 0.8390, 0.2820]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
          UpperCAmelCase_							:   Tuple   =      torch.tensor([-0.2822, -0.0502, -0.0878]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
          UpperCAmelCase_							:   Union[str, Any]   =      torch.tensor([-0.5672, -0.0730, -0.4348]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
          UpperCAmelCase_							:   Optional[Any]   =      torch.tensor([0.2681, 0.2365, 0.6246]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
          UpperCAmelCase_							:   Optional[int]   =      torch.tensor([-0.2642, 0.3931, 0.5116]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
          UpperCAmelCase_							:   Union[str, Any]   =      torch.tensor([-0.6677, -0.1873, -0.8379]   )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
          UpperCAmelCase_							:   str   =      torch.tensor([-0.7749, -0.2967, -0.6444]   )
     else:
          raise ValueError(F'''Unknown URL: {checkpoint_url}'''   )

     assert torch.allclose(logits[0, :3]	,     lowerCamelCase_	,     atol=1e-3   )
     assert logits.shape == expected_shape

     Path(lowerCamelCase_   ).mkdir(exist_ok=lowerCamelCase_   )
     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_   )

     print('Pushing model to the hub...'   )
     UpperCAmelCase_							:   Union[str, Any]   =      'convnext'
     if "tiny" in checkpoint_url:
          model_name += "-tiny"
     elif "small" in checkpoint_url:
          model_name += "-small"
     elif "base" in checkpoint_url:
          model_name += "-base"
     elif "xlarge" in checkpoint_url:
          model_name += "-xlarge"
     elif "large" in checkpoint_url:
          model_name += "-large"
     if "224" in checkpoint_url:
          model_name += "-224"
     elif "384" in checkpoint_url:
          model_name += "-384"
     if "22k" in checkpoint_url and "1k" not in checkpoint_url:
          model_name += "-22k"
     if "22k" in checkpoint_url and "1k" in checkpoint_url:
          model_name += "-22k-1k"

     model.push_to_hub(
         repo_path_or_name=Path(lowerCamelCase_	,     lowerCamelCase_   )	,     organization='nielsr'	,     commit_message='Add model'	,     )


if __name__ == "__main__":
 snake_case__     :			List[Any]      			=  argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--checkpoint_url''',
     default='''https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth''',
     type=str,
     help='''URL of the original ConvNeXT checkpoint you\'d like to convert.''',
 )
 parser.add_argument(
     '''--pytorch_dump_folder_path''',
     default=None,
     type=str,
     required=True,
     help='''Path to the output PyTorch model directory.''',
 )

 snake_case__     :			int      			=  parser.parse_args()
 convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






