






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__ (       __lowerCamelCase     :	Tuple						):
  __UpperCAmelCase				:      str	 =    [2, 2, 6, 2] if """tiny""" in model_name else [2, 2, 18, 2]
  __UpperCAmelCase				:      Any	 =    True if """large""" in model_name or """huge""" in model_name else False
  __UpperCAmelCase				:      int	 =    True if """large""" in model_name or """huge""" in model_name else False
  __UpperCAmelCase				:      Optional[int]	 =    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				:      Union[str, Any]	 =    [3, 3, 3, 3]
      __UpperCAmelCase				:      Union[str, Any]	 =    [5, 5, 5, 5]
    elif "fl4" in model_name:
      __UpperCAmelCase				:      str	 =    [4, 4, 4, 4]
      __UpperCAmelCase				:      Optional[Any]	 =    [3, 3, 3, 3]

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

  if "tiny" in model_name:
    __UpperCAmelCase				:      List[str]	 =    96
  elif "small" in model_name:
    __UpperCAmelCase				:      Dict	 =    96
  elif "base" in model_name:
    __UpperCAmelCase				:      List[Any]	 =    128
  elif "large" in model_name:
    __UpperCAmelCase				:      Any	 =    192
  elif "xlarge" in model_name:
    __UpperCAmelCase				:      Tuple	 =    256
  elif "huge" in model_name:
    __UpperCAmelCase				:      int	 =    352

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

  __UpperCAmelCase				:      str	 =    json.load(open(hf_hub_download(__lowerCamelCase    ,						__lowerCamelCase    ,						repo_type="""dataset"""						)    ,						"""r"""						)						)
  __UpperCAmelCase				:      Optional[int]	 =    {int(__lowerCamelCase						): v for k, v in idalabel.items()}
  __UpperCAmelCase				:      Union[str, Any]	 =    {v: k for k, v in idalabel.items()}

  __UpperCAmelCase				:      List[str]	 =    FocalNetConfig(
      embed_dim=__lowerCamelCase    ,						depths=__lowerCamelCase    ,						focal_levels=__lowerCamelCase    ,						focal_windows=__lowerCamelCase    ,						use_conv_embed=__lowerCamelCase    ,						idalabel=__lowerCamelCase    ,						labelaid=__lowerCamelCase    ,						use_post_layernorm=__lowerCamelCase    ,						use_layerscale=__lowerCamelCase    ,						)

  return config






def 							lowerCamelCase__ (       __lowerCamelCase     :	Tuple						):
  if "patch_embed.proj" in name:
    __UpperCAmelCase				:      List[Any]	 =    name.replace("""patch_embed.proj"""    ,						"""embeddings.patch_embeddings.projection"""						)
  if "patch_embed.norm" in name:
    __UpperCAmelCase				:      Dict	 =    name.replace("""patch_embed.norm"""    ,						"""embeddings.norm"""						)
  if "layers" in name:
    __UpperCAmelCase				:      int	 =    """encoder.""" + name
  if "encoder.layers" in name:
    __UpperCAmelCase				:      Optional[int]	 =    name.replace("""encoder.layers"""    ,						"""encoder.stages"""						)
  if "downsample.proj" in name:
    __UpperCAmelCase				:      Optional[Any]	 =    name.replace("""downsample.proj"""    ,						"""downsample.projection"""						)
  if "blocks" in name:
    __UpperCAmelCase				:      Union[str, Any]	 =    name.replace("""blocks"""    ,						"""layers"""						)
  if "modulation.f.weight" in name or "modulation.f.bias" in name:
    __UpperCAmelCase				:      List[str]	 =    name.replace("""modulation.f"""    ,						"""modulation.projection_in"""						)
  if "modulation.h.weight" in name or "modulation.h.bias" in name:
    __UpperCAmelCase				:      List[Any]	 =    name.replace("""modulation.h"""    ,						"""modulation.projection_context"""						)
  if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
    __UpperCAmelCase				:      str	 =    name.replace("""modulation.proj"""    ,						"""modulation.projection_out"""						)

  if name == "norm.weight":
    __UpperCAmelCase				:      Optional[Any]	 =    """layernorm.weight"""
  if name == "norm.bias":
    __UpperCAmelCase				:      Dict	 =    """layernorm.bias"""

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

  return name






def 							lowerCamelCase__ (       __lowerCamelCase     :	Optional[int]    ,						__lowerCamelCase     :	Optional[Any]    ,						__lowerCamelCase     :	List[str]=False						):
  # fmt: off
  __UpperCAmelCase				:      Dict	 =    {
      """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				:      int	 =    model_name_to_url[model_name]
  print("""Checkpoint URL: """    ,						__lowerCamelCase						)
  __UpperCAmelCase				:      Dict	 =    torch.hub.load_state_dict_from_url(__lowerCamelCase    ,						map_location="""cpu"""						)["""model"""]

  # rename keys
  for key in state_dict.copy().keys():
    __UpperCAmelCase				:      Tuple	 =    state_dict.pop(__lowerCamelCase						)
    __UpperCAmelCase				:      Optional[Any]	 =    val

  __UpperCAmelCase				:      Optional[Any]	 =    get_focalnet_config(__lowerCamelCase						)
  __UpperCAmelCase				:      Any	 =    FocalNetForImageClassification(__lowerCamelCase						)
  model.eval()

  # load state dict
  model.load_state_dict(__lowerCamelCase						)

  # verify conversion
  __UpperCAmelCase				:      Optional[int]	 =    """http://images.cocodataset.org/val2017/000000039769.jpg"""

  __UpperCAmelCase				:      Union[str, Any]	 =    BitImageProcessor(
      do_resize=__lowerCamelCase    ,						size={"""shortest_edge""": 256}    ,						resample=PILImageResampling.BILINEAR    ,						do_center_crop=__lowerCamelCase    ,						crop_size=224    ,						do_normalize=__lowerCamelCase    ,						image_mean=__lowerCamelCase    ,						image_std=__lowerCamelCase    ,						)
  __UpperCAmelCase				:      Optional[Any]	 =    Image.open(requests.get(__lowerCamelCase    ,						stream=__lowerCamelCase						).raw						)
  __UpperCAmelCase				:      List[Any]	 =    processor(images=__lowerCamelCase    ,						return_tensors="""pt"""						)

  __UpperCAmelCase				:      str	 =    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]						),
      ]						)

  __UpperCAmelCase				:      List[Any]	 =    image_transforms(__lowerCamelCase						).unsqueeze(0						)

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

  __UpperCAmelCase				:      Dict	 =    model(**__lowerCamelCase						)

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

  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						:			Union[str, 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.",
       )

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





"""simple docstring"""

import itertools
import os
import random
import tempfile
import unittest

import numpy as np

from transformers import TvltFeatureExtractor, is_datasets_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
    import torch

if is_datasets_available():
    from datasets import load_dataset

_UpperCamelCase		     =    random.Random()



def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__=1.0					,					lowercase__=None					,					lowercase__=None   )			->						Optional[Any]:
 if rng is None:
  lowerCAmelCase__   :				Union[str, Any] 			=				global_rng

 lowerCAmelCase__   :				Dict 			=				[]
 for batch_idx in range(shape[0]   ):
  values.append([]   )
  for _ in range(shape[1]   ):
   values[-1].append(rng.random() * scale   )

 return values






class 			__a					(						unittest.TestCase       ):





       """simple docstring"""


       def __init__(   self			,					snake_case			,					snake_case=7			,					snake_case=400			,					snake_case=2_000			,					snake_case=2_048			,					snake_case=128			,					snake_case=1			,					snake_case=512			,					snake_case=30			,					snake_case=44_100			,					):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				parent
        lowerCAmelCase__   :				Optional[int] 			=				batch_size
        lowerCAmelCase__   :				Optional[Any] 			=				min_seq_length
        lowerCAmelCase__   :				Optional[int] 			=				max_seq_length
        lowerCAmelCase__   :				Dict 			=				(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
        lowerCAmelCase__   :				Dict 			=				spectrogram_length
        lowerCAmelCase__   :				Any 			=				feature_size
        lowerCAmelCase__   :				int 			=				num_audio_channels
        lowerCAmelCase__   :				Optional[int] 			=				hop_length
        lowerCAmelCase__   :				List[str] 			=				chunk_length
        lowerCAmelCase__   :				Optional[Any] 			=				sampling_rate






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        return {
            "spectrogram_length": self.spectrogram_length,
            "feature_size": self.feature_size,
            "num_audio_channels": self.num_audio_channels,
            "hop_length": self.hop_length,
            "chunk_length": self.chunk_length,
            "sampling_rate": self.sampling_rate,
        }






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case=False			,					snake_case=False   ):





        """simple docstring"""

        def _flatten(snake_case   ):
         return list(itertools.chain(*snake_case   )   )

        if equal_length:
         lowerCAmelCase__   :				Dict 			=				[floats_list((self.max_seq_length, self.feature_size)   ) for _ in range(self.batch_size   )]
        else:
         # make sure that inputs increase in size
         lowerCAmelCase__   :				Tuple 			=				[
             floats_list((x, self.feature_size)   )
             for x in range(self.min_seq_length			,					self.max_seq_length			,					self.seq_length_diff   )
         ]
        if numpify:
         lowerCAmelCase__   :				Dict 			=				[np.asarray(snake_case   ) for x in speech_inputs]
        return speech_inputs






@require_torch
@require_torchaudio
class 			__a					(						__magic_name__       ,	unittest.TestCase       ):





       """simple docstring"""


       __UpperCamelCase					:      Optional[Any]    =						TvltFeatureExtractor






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[Any] 			=				TvltFeatureExtractionTester(self   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[str] 			=				self.feature_extraction_class(**self.feat_extract_dict   )
        self.assertTrue(hasattr(snake_case			,					"spectrogram_length"   )   )
        self.assertTrue(hasattr(snake_case			,					"feature_size"   )   )
        self.assertTrue(hasattr(snake_case			,					"num_audio_channels"   )   )
        self.assertTrue(hasattr(snake_case			,					"hop_length"   )   )
        self.assertTrue(hasattr(snake_case			,					"chunk_length"   )   )
        self.assertTrue(hasattr(snake_case			,					"sampling_rate"   )   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				int 			=				self.feature_extraction_class(**self.feat_extract_dict   )

        with tempfile.TemporaryDirectory() as tmpdirname:
         lowerCAmelCase__   :				Union[str, Any] 			=				feat_extract_first.save_pretrained(snake_case   )[0]
         check_json_file_has_correct_format(snake_case   )
         lowerCAmelCase__   :				Tuple 			=				self.feature_extraction_class.from_pretrained(snake_case   )

        lowerCAmelCase__   :				Union[str, Any] 			=				feat_extract_first.to_dict()
        lowerCAmelCase__   :				int 			=				feat_extract_second.to_dict()
        lowerCAmelCase__   :				Union[str, Any] 			=				dict_first.pop("mel_filters"   )
        lowerCAmelCase__   :				List[Any] 			=				dict_second.pop("mel_filters"   )
        self.assertTrue(np.allclose(snake_case			,					snake_case   )   )
        self.assertEqual(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Any 			=				self.feature_extraction_class(**self.feat_extract_dict   )

        with tempfile.TemporaryDirectory() as tmpdirname:
         lowerCAmelCase__   :				List[Any] 			=				os.path.join(snake_case			,					"feat_extract.json"   )
         feat_extract_first.to_json_file(snake_case   )
         lowerCAmelCase__   :				List[Any] 			=				self.feature_extraction_class.from_json_file(snake_case   )

        lowerCAmelCase__   :				Optional[Any] 			=				feat_extract_first.to_dict()
        lowerCAmelCase__   :				Dict 			=				feat_extract_second.to_dict()
        lowerCAmelCase__   :				Tuple 			=				dict_first.pop("mel_filters"   )
        lowerCAmelCase__   :				Dict 			=				dict_second.pop("mel_filters"   )
        self.assertTrue(np.allclose(snake_case			,					snake_case   )   )
        self.assertEqual(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				self.feature_extraction_class(**self.feat_extract_dict   )

        # create three inputs of length 800, 1000, and 1200
        lowerCAmelCase__   :				Union[str, Any] 			=				[floats_list((1, x)   )[0] for x in range(800			,					1_400			,					200   )]
        lowerCAmelCase__   :				Optional[int] 			=				[np.asarray(snake_case   ) for speech_input in speech_inputs]

        # Test not batched input
        lowerCAmelCase__   :				Tuple 			=				feature_extractor(np_speech_inputs[0]			,					return_tensors="np"			,					sampling_rate=44_100   ).audio_values

        self.assertTrue(encoded_audios.ndim == 4   )
        self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size   )
        self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length   )
        self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels   )

        # Test batched
        lowerCAmelCase__   :				Optional[int] 			=				feature_extractor(snake_case			,					return_tensors="np"			,					sampling_rate=44_100   ).audio_values

        self.assertTrue(encoded_audios.ndim == 4   )
        self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size   )
        self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length   )
        self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels   )

        # Test audio masking
        lowerCAmelCase__   :				List[str] 			=				feature_extractor(
            snake_case			,					return_tensors="np"			,					sampling_rate=44_100			,					mask_audio=snake_case   ).audio_values

        self.assertTrue(encoded_audios.ndim == 4   )
        self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size   )
        self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length   )
        self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels   )

        # Test 2-D numpy arrays are batched.
        lowerCAmelCase__   :				List[str] 			=				[floats_list((1, x)   )[0] for x in (800, 800, 800)]
        lowerCAmelCase__   :				Tuple 			=				np.asarray(snake_case   )
        lowerCAmelCase__   :				List[Any] 			=				feature_extractor(snake_case			,					return_tensors="np"			,					sampling_rate=44_100   ).audio_values
        self.assertTrue(encoded_audios.ndim == 4   )
        self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size   )
        self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length   )
        self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels   )






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				str 			=				load_dataset("hf-internal-testing/librispeech_asr_dummy"			,					"clean"			,					split="validation"   )
        # automatic decoding with librispeech
        lowerCAmelCase__   :				int 			=				ds.sort("id"   ).select(range(snake_case   )   )[:num_samples]["audio"]

        return [x["array"] for x in speech_samples]






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				str 			=				self._load_datasamples(1   )
        lowerCAmelCase__   :				Optional[Any] 			=				TvltFeatureExtractor()
        lowerCAmelCase__   :				Dict 			=				feature_extractor(snake_case			,					return_tensors="pt"   ).audio_values

        self.assertEquals(audio_values.shape			,					(1, 1, 192, 128)   )

        lowerCAmelCase__   :				Dict 			=				torch.tensor([[-0.3_032, -0.2_708], [-0.4_434, -0.4_007]]   )
        self.assertTrue(torch.allclose(audio_values[0, 0, :2, :2]			,					snake_case			,					atol=1e-4   )   )






