




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 YolosConfig, YolosForObjectDetection, YolosImageProcessor
from transformers.utils import logging


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







def 	_a      (						SCREAMING_SNAKE_CASE_							:				str     ):
					__lowerCAmelCase				   =       YolosConfig()

					# size of the architecture
					if "yolos_ti" in yolos_name:
										__lowerCAmelCase				   =       1_92
										__lowerCAmelCase				   =       7_68
										__lowerCAmelCase				   =       12
										__lowerCAmelCase				   =       3
										__lowerCAmelCase				   =       [8_00, 13_33]
										__lowerCAmelCase				   =       False
					elif yolos_name == "yolos_s_dWr":
										__lowerCAmelCase				   =       3_30
										__lowerCAmelCase				   =       14
										__lowerCAmelCase				   =       6
										__lowerCAmelCase				   =       13_20
					elif "yolos_s" in yolos_name:
										__lowerCAmelCase				   =       3_84
										__lowerCAmelCase				   =       15_36
										__lowerCAmelCase				   =       12
										__lowerCAmelCase				   =       6
					elif "yolos_b" in yolos_name:
										__lowerCAmelCase				   =       [8_00, 13_44]

					__lowerCAmelCase				   =       91
					__lowerCAmelCase				   =       "huggingface/label-files"
					__lowerCAmelCase				   =       "coco-detection-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()}

					return config







def 	_a      (						SCREAMING_SNAKE_CASE_							:				dict      ,      SCREAMING_SNAKE_CASE_							:				YolosConfig      ,      SCREAMING_SNAKE_CASE_							:				bool = False     ):
					for i in range(config.num_hidden_layers     ):
										# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
										__lowerCAmelCase				   =       state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""     )
										__lowerCAmelCase				   =       state_dict.pop(F"""blocks.{i}.attn.qkv.bias"""     )
										# next, add query, keys and values (in that order) to the state dict
										__lowerCAmelCase				   =       in_proj_weight[: config.hidden_size, :]
										__lowerCAmelCase				   =       in_proj_bias[: config.hidden_size]
										__lowerCAmelCase				   =       in_proj_weight[
										    config.hidden_size : config.hidden_size * 2, :
										]
										__lowerCAmelCase				   =       in_proj_bias[
										    config.hidden_size : config.hidden_size * 2
										]
										__lowerCAmelCase				   =       in_proj_weight[-config.hidden_size :, :]
										__lowerCAmelCase				   =       in_proj_bias[-config.hidden_size :]







def 	_a      (						SCREAMING_SNAKE_CASE_							:				str     ):
					if "backbone" in name:
										__lowerCAmelCase				   =       name.replace("backbone"      ,      "vit"     )
					if "cls_token" in name:
										__lowerCAmelCase				   =       name.replace("cls_token"      ,      "embeddings.cls_token"     )
					if "det_token" in name:
										__lowerCAmelCase				   =       name.replace("det_token"      ,      "embeddings.detection_tokens"     )
					if "mid_pos_embed" in name:
										__lowerCAmelCase				   =       name.replace("mid_pos_embed"      ,      "encoder.mid_position_embeddings"     )
					if "pos_embed" in name:
										__lowerCAmelCase				   =       name.replace("pos_embed"      ,      "embeddings.position_embeddings"     )
					if "patch_embed.proj" in name:
										__lowerCAmelCase				   =       name.replace("patch_embed.proj"      ,      "embeddings.patch_embeddings.projection"     )
					if "blocks" in name:
										__lowerCAmelCase				   =       name.replace("blocks"      ,      "encoder.layer"     )
					if "attn.proj" in name:
										__lowerCAmelCase				   =       name.replace("attn.proj"      ,      "attention.output.dense"     )
					if "attn" in name:
										__lowerCAmelCase				   =       name.replace("attn"      ,      "attention.self"     )
					if "norm1" in name:
										__lowerCAmelCase				   =       name.replace("norm1"      ,      "layernorm_before"     )
					if "norm2" in name:
										__lowerCAmelCase				   =       name.replace("norm2"      ,      "layernorm_after"     )
					if "mlp.fc1" in name:
										__lowerCAmelCase				   =       name.replace("mlp.fc1"      ,      "intermediate.dense"     )
					if "mlp.fc2" in name:
										__lowerCAmelCase				   =       name.replace("mlp.fc2"      ,      "output.dense"     )
					if "class_embed" in name:
										__lowerCAmelCase				   =       name.replace("class_embed"      ,      "class_labels_classifier"     )
					if "bbox_embed" in name:
										__lowerCAmelCase				   =       name.replace("bbox_embed"      ,      "bbox_predictor"     )
					if "vit.norm" in name:
										__lowerCAmelCase				   =       name.replace("vit.norm"      ,      "vit.layernorm"     )

					return name







def 	_a      (						SCREAMING_SNAKE_CASE_							:				dict      ,      SCREAMING_SNAKE_CASE_							:				YolosForObjectDetection     ):
					for key in orig_state_dict.copy().keys():
										__lowerCAmelCase				   =       orig_state_dict.pop(SCREAMING_SNAKE_CASE_     )

										if "qkv" in key:
															__lowerCAmelCase				   =       key.split("."     )
															__lowerCAmelCase				   =       int(key_split[2]     )
															__lowerCAmelCase				   =       model.vit.encoder.layer[layer_num].attention.attention.all_head_size
															if "weight" in key:
																				__lowerCAmelCase				   =       val[:dim, :]
																				__lowerCAmelCase				   =       val[
																				    dim : dim * 2, :
																				]
																				__lowerCAmelCase				   =       val[-dim:, :]
															else:
																				__lowerCAmelCase				   =       val[:dim]
																				__lowerCAmelCase				   =       val[dim : dim * 2]
																				__lowerCAmelCase				   =       val[-dim:]
										else:
															__lowerCAmelCase				   =       val

					return orig_state_dict







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







@torch.no_grad()
def 	_a      (						SCREAMING_SNAKE_CASE_							:				str      ,      SCREAMING_SNAKE_CASE_							:				str      ,      SCREAMING_SNAKE_CASE_							:				str      ,      SCREAMING_SNAKE_CASE_							:				bool = False     ):
					__lowerCAmelCase				   =       get_yolos_config(SCREAMING_SNAKE_CASE_     )

					# load original state_dict
					__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_      ,      map_location="cpu"     )["model"]

					# load 🤗 model
					__lowerCAmelCase				   =       YolosForObjectDetection(SCREAMING_SNAKE_CASE_     )
					model.eval()
					__lowerCAmelCase				   =       convert_state_dict(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
					model.load_state_dict(SCREAMING_SNAKE_CASE_     )

					# Check outputs on an image, prepared by YolosImageProcessor
					__lowerCAmelCase				   =       8_00 if yolos_name != "yolos_ti" else 5_12
					__lowerCAmelCase				   =       YolosImageProcessor(format="coco_detection"      ,      size=SCREAMING_SNAKE_CASE_     )
					__lowerCAmelCase				   =       image_processor(images=prepare_img()      ,      return_tensors="pt"     )
					__lowerCAmelCase				   =       model(**SCREAMING_SNAKE_CASE_     )
					__lowerCAmelCase   ,							__lowerCAmelCase				   =       outputs.logits, outputs.pred_boxes

					__lowerCAmelCase   ,							__lowerCAmelCase				   =       None, None
					if yolos_name == "yolos_ti":
										__lowerCAmelCase				   =       torch.tensor(
										    [[-39.50_22, -11.98_20, -17.68_88], [-29.95_74, -9.97_69, -17.76_91], [-42.32_81, -20.72_00, -30.62_94]]     )
										__lowerCAmelCase				   =       torch.tensor(
										    [[0.40_21, 0.08_36, 0.79_79], [0.01_84, 0.26_09, 0.03_64], [0.17_81, 0.20_04, 0.20_95]]     )
					elif yolos_name == "yolos_s_200_pre":
										__lowerCAmelCase				   =       torch.tensor(
										    [[-24.02_48, -10.30_24, -14.82_90], [-42.03_92, -16.82_00, -27.43_34], [-27.27_43, -11.81_54, -18.71_48]]     )
										__lowerCAmelCase				   =       torch.tensor(
										    [[0.25_59, 0.54_55, 0.47_06], [0.29_89, 0.72_79, 0.18_75], [0.77_32, 0.40_17, 0.44_62]]     )
					elif yolos_name == "yolos_s_300_pre":
										__lowerCAmelCase				   =       torch.tensor(
										    [[-36.22_20, -14.43_85, -23.54_57], [-35.69_70, -14.75_83, -21.39_35], [-31.59_39, -13.60_42, -16.80_49]]     )
										__lowerCAmelCase				   =       torch.tensor(
										    [[0.76_14, 0.23_16, 0.47_28], [0.71_68, 0.44_95, 0.38_55], [0.49_96, 0.14_66, 0.99_96]]     )
					elif yolos_name == "yolos_s_dWr":
										__lowerCAmelCase				   =       torch.tensor(
										    [[-42.86_68, -24.10_49, -41.16_90], [-34.74_56, -14.12_74, -24.91_94], [-33.78_98, -12.19_46, -25.64_95]]     )
										__lowerCAmelCase				   =       torch.tensor(
										    [[0.55_87, 0.27_73, 0.06_05], [0.50_04, 0.30_14, 0.99_94], [0.49_99, 0.15_48, 0.99_94]]     )
					elif yolos_name == "yolos_base":
										__lowerCAmelCase				   =       torch.tensor(
										    [[-40.60_64, -24.30_84, -32.64_47], [-55.19_90, -30.77_19, -35.58_77], [-51.43_11, -33.35_07, -35.64_62]]     )
										__lowerCAmelCase				   =       torch.tensor(
										    [[0.55_55, 0.27_94, 0.06_55], [0.90_49, 0.26_64, 0.18_94], [0.91_83, 0.19_84, 0.16_35]]     )
					else:
										raise ValueError(F"""Unknown yolos_name: {yolos_name}"""     )

					assert torch.allclose(logits[0, :3, :3]      ,      SCREAMING_SNAKE_CASE_      ,      atol=1E-4     )
					assert torch.allclose(pred_boxes[0, :3, :3]      ,      SCREAMING_SNAKE_CASE_      ,      atol=1E-4     )

					Path(SCREAMING_SNAKE_CASE_     ).mkdir(exist_ok=SCREAMING_SNAKE_CASE_     )
					print(F"""Saving model {yolos_name} to {pytorch_dump_folder_path}"""     )
					model.save_pretrained(SCREAMING_SNAKE_CASE_     )
					print(F"""Saving image processor to {pytorch_dump_folder_path}"""     )
					image_processor.save_pretrained(SCREAMING_SNAKE_CASE_     )

					if push_to_hub:
										__lowerCAmelCase				   =       {
										    "yolos_ti": "yolos-tiny",
										    "yolos_s_200_pre": "yolos-small",
										    "yolos_s_300_pre": "yolos-small-300",
										    "yolos_s_dWr": "yolos-small-dwr",
										    "yolos_base": "yolos-base",
										}

										print("Pushing to the hub..."     )
										__lowerCAmelCase				   =       model_mapping[yolos_name]
										image_processor.push_to_hub(SCREAMING_SNAKE_CASE_      ,      organization="hustvl"     )
										model.push_to_hub(SCREAMING_SNAKE_CASE_      ,      organization="hustvl"     )


if __name__ == "__main__":
					UpperCamelCase__    			=			argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--yolos_name""",
					    default="""yolos_s_200_pre""",
					    type=str,
					    help=(
					        """Name of the YOLOS model you'd like to convert. Should be one of 'yolos_ti', 'yolos_s_200_pre',"""
					        """ 'yolos_s_300_pre', 'yolos_s_dWr', 'yolos_base'."""
					    ),
					)
					parser.add_argument(
					    """--checkpoint_path""", default=None, type=str, help="""Path to the original state dict (.pth file)."""
					)
					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 or not to push the converted model to the 🤗 hub."""
					)

					UpperCamelCase__    			=			parser.parse_args()
					convert_yolos_checkpoint(args.yolos_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)










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__)




