




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 (
    SwiftFormerConfig,
    SwiftFormerForImageClassification,
    ViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase__    			=			logging.get_logger(__name__)

UpperCamelCase__    			=			torch.device("""cpu""")







def 	_a      (						):
					__lowerCAmelCase				   =       "http://images.cocodataset.org/val2017/000000039769.jpg"
					__lowerCAmelCase				   =       Image.open(requests.get(SCREAMING_SNAKE_CASE_      ,      stream=SCREAMING_SNAKE_CASE_     ).raw     )
					return im







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Any     ):
					if swiftformer_name == "swiftformer_xs":
										return torch.tensor([-2.17_03E00, 2.11_07E00, -2.08_11E00, 8.86_85E-01, 2.43_60E-01]     )

					elif swiftformer_name == "swiftformer_s":
										return torch.tensor([3.96_36E-01, 2.34_78E-01, -1.69_63E00, -1.73_81E00, -8.63_37E-01]     )

					elif swiftformer_name == "swiftformer_l1":
										return torch.tensor([-4.27_68E-01, -4.74_29E-01, -1.08_97E00, -1.02_48E00, 3.55_23E-02]     )

					elif swiftformer_name == "swiftformer_l3":
										return torch.tensor([-2.53_30E-01, 2.42_11E-01, -6.01_85E-01, -8.27_89E-01, -6.04_46E-02]     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Any      ,      SCREAMING_SNAKE_CASE_							:				Any      ,      SCREAMING_SNAKE_CASE_							:				int     ):
					__lowerCAmelCase				   =       dct.pop(SCREAMING_SNAKE_CASE_     )
					__lowerCAmelCase				   =       val







def 	_a      (						SCREAMING_SNAKE_CASE_							:				str     ):
					__lowerCAmelCase				   =       []
					for k in state_dict.keys():
										__lowerCAmelCase				   =       k
										if ".pwconv" in k:
															__lowerCAmelCase				   =       k_new.replace(".pwconv"      ,      ".point_wise_conv"     )
										if ".dwconv" in k:
															__lowerCAmelCase				   =       k_new.replace(".dwconv"      ,      ".depth_wise_conv"     )
										if ".Proj." in k:
															__lowerCAmelCase				   =       k_new.replace(".Proj."      ,      ".proj."     )
										if "patch_embed" in k_new:
															__lowerCAmelCase				   =       k_new.replace("patch_embed"      ,      "swiftformer.patch_embed.patch_embedding"     )
										if "network" in k_new:
															__lowerCAmelCase				   =       k_new.split("."     )
															if ls[2].isdigit():
																				__lowerCAmelCase				   =       "swiftformer.encoder.network." + ls[1] + ".blocks." + ls[2] + "." + ".".join(ls[3:]     )
															else:
																				__lowerCAmelCase				   =       k_new.replace("network"      ,      "swiftformer.encoder.network"     )
										rename_keys.append((k, k_new)     )
					return rename_keys







@torch.no_grad()
def 	_a      (						SCREAMING_SNAKE_CASE_							:				Any      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				Tuple     ):
					__lowerCAmelCase				   =       SwiftFormerConfig()

					# dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
					__lowerCAmelCase				   =       10_00
					__lowerCAmelCase				   =       "huggingface/label-files"
					__lowerCAmelCase				   =       "imagenet-1k-id2label.json"
					__lowerCAmelCase				   =       json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      repo_type="dataset"     )      ,      "r"     )     )
					__lowerCAmelCase				   =       {int(SCREAMING_SNAKE_CASE_     ): v for k, v in idalabel.items()}
					__lowerCAmelCase				   =       idalabel
					__lowerCAmelCase				   =       {v: k for k, v in idalabel.items()}

					# size of the architecture
					if swiftformer_name == "swiftformer_xs":
										__lowerCAmelCase				   =       [3, 3, 6, 4]
										__lowerCAmelCase				   =       [48, 56, 1_12, 2_20]

					elif swiftformer_name == "swiftformer_s":
										__lowerCAmelCase				   =       [3, 3, 9, 6]
										__lowerCAmelCase				   =       [48, 64, 1_68, 2_24]

					elif swiftformer_name == "swiftformer_l1":
										__lowerCAmelCase				   =       [4, 3, 10, 5]
										__lowerCAmelCase				   =       [48, 96, 1_92, 3_84]

					elif swiftformer_name == "swiftformer_l3":
										__lowerCAmelCase				   =       [4, 4, 12, 6]
										__lowerCAmelCase				   =       [64, 1_28, 3_20, 5_12]

					# load state_dict of original model, remove and rename some keys
					if original_ckpt:
										if original_ckpt.startswith("https"     ):
															__lowerCAmelCase				   =       torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_      ,      map_location="cpu"      ,      check_hash=SCREAMING_SNAKE_CASE_     )
										else:
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_      ,      map_location="cpu"     )
					__lowerCAmelCase				   =       checkpoint

					__lowerCAmelCase				   =       create_rename_keys(SCREAMING_SNAKE_CASE_     )
					for rename_key_src, rename_key_dest in rename_keys:
										rename_key(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )

					# load HuggingFace model
					__lowerCAmelCase				   =       SwiftFormerForImageClassification(SCREAMING_SNAKE_CASE_     ).eval()
					hf_model.load_state_dict(SCREAMING_SNAKE_CASE_     )

					# prepare test inputs
					__lowerCAmelCase				   =       prepare_img()
					__lowerCAmelCase				   =       ViTImageProcessor.from_pretrained("preprocessor_config"     )
					__lowerCAmelCase				   =       processor(images=SCREAMING_SNAKE_CASE_      ,      return_tensors="pt"     )

					# compare outputs from both models
					__lowerCAmelCase				   =       get_expected_output(SCREAMING_SNAKE_CASE_     )
					__lowerCAmelCase				   =       hf_model(inputs["pixel_values"]     ).logits

					assert hf_logits.shape == torch.Size([1, 10_00]     )
					assert torch.allclose(hf_logits[0, 0:5]      ,      SCREAMING_SNAKE_CASE_      ,      atol=1E-3     )

					Path(SCREAMING_SNAKE_CASE_     ).mkdir(exist_ok=SCREAMING_SNAKE_CASE_     )
					print(F"""Saving model {swiftformer_name} to {pytorch_dump_folder_path}"""     )
					hf_model.save_pretrained(SCREAMING_SNAKE_CASE_     )


if __name__ == "__main__":
					UpperCamelCase__    			=			argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--swiftformer_name""",
					    default="""swiftformer_xs""",
					    choices=["""swiftformer_xs""", """swiftformer_s""", """swiftformer_l1""", """swiftformer_l3"""],
					    type=str,
					    help="""Name of the SwiftFormer model you'd like to convert.""",
					)
					parser.add_argument(
					    """--pytorch_dump_folder_path""",
					    default="""./converted_outputs/""",
					    type=str,
					    help="""Path to the output PyTorch model directory.""",
					)
					parser.add_argument("""--original_ckpt""", default=None, type=str, help="""Path to the original model checkpoint.""")

					UpperCamelCase__    			=			parser.parse_args()
					convert_swiftformer_checkpoint(args.swiftformer_name, args.pytorch_dump_folder_path, args.original_ckpt)










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_speech_available, is_torch_available


UpperCamelCase__    			=			{
    """configuration_audio_spectrogram_transformer""": [
        """AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """ASTConfig""",
    ]
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCamelCase__    			=			[
					    """AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """ASTForAudioClassification""",
					    """ASTModel""",
					    """ASTPreTrainedModel""",
					]

try:
					if not is_speech_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCamelCase__    			=			["""ASTFeatureExtractor"""]

if TYPE_CHECKING:
					from .configuration_audio_spectrogram_transformer import (
					    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
					    ASTConfig,
					)

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_audio_spectrogram_transformer import (
										    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    ASTForAudioClassification,
										    ASTModel,
										    ASTPreTrainedModel,
										)

					try:
										if not is_speech_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .feature_extraction_audio_spectrogram_transformer import ASTFeatureExtractor


else:
					import sys

					UpperCamelCase__    			=			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




