

'''simple docstring'''






import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
from PIL import Image

from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor, SwinConfig
from transformers.utils import logging


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


def 	_A  (					A__   ):






 """simple docstring"""



 __lowercase													=				SwinConfig(
     embed_dim=192  ,       depths=(2, 2, 18, 2)  ,       num_heads=(6, 12, 24, 48)  ,       window_size=12  ,       out_features=['''stage2''', '''stage3''', '''stage4''']  ,       )

 __lowercase													=				DetaConfig(
     backbone_config=A__  ,       num_queries=900  ,       encoder_ffn_dim=2048  ,       decoder_ffn_dim=2048  ,       num_feature_levels=5  ,       assign_first_stage=A__  ,       with_box_refine=A__  ,       two_stage=A__  ,       )

 # set labels
 __lowercase													=				'''huggingface/label-files'''
 if "o365" in model_name:
  __lowercase													=				366
  __lowercase													=				'''object365-id2label.json'''
 else:
  __lowercase													=				91
  __lowercase													=				'''coco-detection-id2label.json'''

 __lowercase													=				num_labels
 __lowercase													=				json.load(open(cached_download(hf_hub_url(A__  ,       A__  ,       repo_type='''dataset'''   )   )  ,       '''r'''   )   )
 __lowercase													=				{int(A__   ): v for k, v in idalabel.items()}
 __lowercase													=				idalabel
 __lowercase													=				{v: k for k, v in idalabel.items()}

 return config


def 	_A  (					A__   ):






 """simple docstring"""



 __lowercase													=				[]

 # stem
 # fmt: off
 rename_keys.append(('''backbone.0.body.patch_embed.proj.weight''', '''model.backbone.model.embeddings.patch_embeddings.projection.weight''')   )
 rename_keys.append(('''backbone.0.body.patch_embed.proj.bias''', '''model.backbone.model.embeddings.patch_embeddings.projection.bias''')   )
 rename_keys.append(('''backbone.0.body.patch_embed.norm.weight''', '''model.backbone.model.embeddings.norm.weight''')   )
 rename_keys.append(('''backbone.0.body.patch_embed.norm.bias''', '''model.backbone.model.embeddings.norm.bias''')   )
 # stages
 for i in range(len(config.backbone_config.depths   )   ):
  for j in range(config.backbone_config.depths[i]   ):
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.norm1.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.norm1.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_before.bias")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_bias_table", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.attn.relative_position_index", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.attn.proj.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.attn.proj.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.norm2.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.norm2.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.layernorm_after.bias")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc1.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.weight", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.blocks.{j}.mlp.fc2.bias", F"model.backbone.model.encoder.layers.{i}.blocks.{j}.output.dense.bias")   )

  if i < 3:
   rename_keys.append((F"backbone.0.body.layers.{i}.downsample.reduction.weight", F"model.backbone.model.encoder.layers.{i}.downsample.reduction.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.downsample.norm.weight", F"model.backbone.model.encoder.layers.{i}.downsample.norm.weight")   )
   rename_keys.append((F"backbone.0.body.layers.{i}.downsample.norm.bias", F"model.backbone.model.encoder.layers.{i}.downsample.norm.bias")   )

 rename_keys.append(('''backbone.0.body.norm1.weight''', '''model.backbone.model.hidden_states_norms.stage2.weight''')   )
 rename_keys.append(('''backbone.0.body.norm1.bias''', '''model.backbone.model.hidden_states_norms.stage2.bias''')   )
 rename_keys.append(('''backbone.0.body.norm2.weight''', '''model.backbone.model.hidden_states_norms.stage3.weight''')   )
 rename_keys.append(('''backbone.0.body.norm2.bias''', '''model.backbone.model.hidden_states_norms.stage3.bias''')   )
 rename_keys.append(('''backbone.0.body.norm3.weight''', '''model.backbone.model.hidden_states_norms.stage4.weight''')   )
 rename_keys.append(('''backbone.0.body.norm3.bias''', '''model.backbone.model.hidden_states_norms.stage4.bias''')   )

 # transformer encoder
 for i in range(config.encoder_layers   ):
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight", F"model.encoder.layers.{i}.self_attn.sampling_offsets.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias", F"model.encoder.layers.{i}.self_attn.sampling_offsets.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.attention_weights.weight", F"model.encoder.layers.{i}.self_attn.attention_weights.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.attention_weights.bias", F"model.encoder.layers.{i}.self_attn.attention_weights.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.value_proj.weight", F"model.encoder.layers.{i}.self_attn.value_proj.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.value_proj.bias", F"model.encoder.layers.{i}.self_attn.value_proj.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.output_proj.weight", F"model.encoder.layers.{i}.self_attn.output_proj.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.self_attn.output_proj.bias", F"model.encoder.layers.{i}.self_attn.output_proj.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.norm1.weight", F"model.encoder.layers.{i}.self_attn_layer_norm.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.norm1.bias", F"model.encoder.layers.{i}.self_attn_layer_norm.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.linear1.weight", F"model.encoder.layers.{i}.fc1.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.linear1.bias", F"model.encoder.layers.{i}.fc1.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.linear2.weight", F"model.encoder.layers.{i}.fc2.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.linear2.bias", F"model.encoder.layers.{i}.fc2.bias")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.norm2.weight", F"model.encoder.layers.{i}.final_layer_norm.weight")   )
  rename_keys.append((F"transformer.encoder.layers.{i}.norm2.bias", F"model.encoder.layers.{i}.final_layer_norm.bias")   )

 # transformer decoder
 for i in range(config.decoder_layers   ):
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight", F"model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias", F"model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.attention_weights.weight", F"model.decoder.layers.{i}.encoder_attn.attention_weights.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.attention_weights.bias", F"model.decoder.layers.{i}.encoder_attn.attention_weights.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.value_proj.weight", F"model.decoder.layers.{i}.encoder_attn.value_proj.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.value_proj.bias", F"model.decoder.layers.{i}.encoder_attn.value_proj.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.output_proj.weight", F"model.decoder.layers.{i}.encoder_attn.output_proj.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.cross_attn.output_proj.bias", F"model.decoder.layers.{i}.encoder_attn.output_proj.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm1.weight", F"model.decoder.layers.{i}.encoder_attn_layer_norm.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm1.bias", F"model.decoder.layers.{i}.encoder_attn_layer_norm.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.self_attn.out_proj.weight", F"model.decoder.layers.{i}.self_attn.out_proj.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.self_attn.out_proj.bias", F"model.decoder.layers.{i}.self_attn.out_proj.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm2.weight", F"model.decoder.layers.{i}.self_attn_layer_norm.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm2.bias", F"model.decoder.layers.{i}.self_attn_layer_norm.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.linear1.weight", F"model.decoder.layers.{i}.fc1.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.linear1.bias", F"model.decoder.layers.{i}.fc1.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.linear2.weight", F"model.decoder.layers.{i}.fc2.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.linear2.bias", F"model.decoder.layers.{i}.fc2.bias")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm3.weight", F"model.decoder.layers.{i}.final_layer_norm.weight")   )
  rename_keys.append((F"transformer.decoder.layers.{i}.norm3.bias", F"model.decoder.layers.{i}.final_layer_norm.bias")   )

 # fmt: on

 return rename_keys


def 	_A  (					A__  ,       A__  ,       A__   ):






 """simple docstring"""



 __lowercase													=				dct.pop(A__   )
 __lowercase													=				val


def 	_A  (					A__  ,       A__   ):






 """simple docstring"""



 __lowercase													=				[int(backbone_config.embed_dim * 2**i   ) for i in range(len(backbone_config.depths   )   )]
 for i in range(len(backbone_config.depths   )   ):
  __lowercase													=				num_features[i]
  for j in range(backbone_config.depths[i]   ):
   # fmt: off
   # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
   __lowercase													=				state_dict.pop(F"backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight"   )
   __lowercase													=				state_dict.pop(F"backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.bias"   )
   # next, add query, keys and values (in that order) to the state dict
   __lowercase													=				in_proj_weight[:dim, :]
   __lowercase													=				in_proj_bias[: dim]
   __lowercase													=				in_proj_weight[
       dim : dim * 2, :
   ]
   __lowercase													=				in_proj_bias[
       dim : dim * 2
   ]
   __lowercase													=				in_proj_weight[
       -dim :, :
   ]
   __lowercase													=				in_proj_bias[-dim :]
   # fmt: on


def 	_A  (					A__  ,       A__   ):






 """simple docstring"""



 __lowercase													=				config.d_model
 for i in range(config.decoder_layers   ):
  # read in weights + bias of input projection layer of self-attention
  __lowercase													=				state_dict.pop(F"transformer.decoder.layers.{i}.self_attn.in_proj_weight"   )
  __lowercase													=				state_dict.pop(F"transformer.decoder.layers.{i}.self_attn.in_proj_bias"   )
  # next, add query, keys and values (in that order) to the state dict
  __lowercase													=				in_proj_weight[:hidden_size, :]
  __lowercase													=				in_proj_bias[:hidden_size]
  __lowercase													=				in_proj_weight[
      hidden_size : hidden_size * 2, :
  ]
  __lowercase													=				in_proj_bias[hidden_size : hidden_size * 2]
  __lowercase													=				in_proj_weight[-hidden_size:, :]
  __lowercase													=				in_proj_bias[-hidden_size:]


def 	_A  (					):






 """simple docstring"""



 __lowercase													=				'''http://images.cocodataset.org/val2017/000000039769.jpg'''
 __lowercase													=				Image.open(requests.get(A__  ,       stream=A__   ).raw   )

 return im


@torch.no_grad()
def 	_A  (					A__  ,       A__  ,       A__   ):






 """simple docstring"""



 __lowercase													=				get_deta_config(A__   )

 # load original state dict
 if model_name == "deta-swin-large":
  __lowercase													=				hf_hub_download(repo_id='''nielsr/deta-checkpoints'''  ,       filename='''adet_swin_ft.pth'''   )
 elif model_name == "deta-swin-large-o365":
  __lowercase													=				hf_hub_download(repo_id='''jozhang97/deta-swin-l-o365'''  ,       filename='''deta_swin_pt_o365.pth'''   )
 else:
  raise ValueError(F"Model name {model_name} not supported"   )

 __lowercase													=				torch.load(A__  ,       map_location='''cpu'''   )['''model''']

 # original state dict
 for name, param in state_dict.items():
  print(A__  ,       param.shape   )

 # rename keys
 __lowercase													=				create_rename_keys(A__   )
 for src, dest in rename_keys:
  rename_key(A__  ,       A__  ,       A__   )
 read_in_swin_q_k_v(A__  ,       config.backbone_config   )
 read_in_decoder_q_k_v(A__  ,       A__   )

 # fix some prefixes
 for key in state_dict.copy().keys():
  if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
   __lowercase													=				state_dict.pop(A__   )
   __lowercase													=				val
  if "input_proj" in key:
   __lowercase													=				state_dict.pop(A__   )
   __lowercase													=				val
  if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
   __lowercase													=				state_dict.pop(A__   )
   __lowercase													=				val

    # finally, create HuggingFace model and load state dict
 __lowercase													=				DetaForObjectDetection(A__   )
 model.load_state_dict(A__   )
 model.eval()

 __lowercase													=				'''cuda''' if torch.cuda.is_available() else '''cpu'''
 model.to(A__   )

 # load image processor
 __lowercase													=				DetaImageProcessor(format='''coco_detection'''   )

 # verify our conversion on image
 __lowercase													=				prepare_img()
 __lowercase													=				processor(images=A__  ,       return_tensors='''pt'''   )
 __lowercase													=				encoding['''pixel_values''']
 __lowercase													=				model(pixel_values.to(A__   )   )

 # verify logits
 print('''Logits:'''  ,       outputs.logits[0, :3, :3]   )
 print('''Boxes:'''  ,       outputs.pred_boxes[0, :3, :3]   )
 if model_name == "deta-swin-large":
  __lowercase													=				torch.tensor(
      [[-7.6_3_0_8, -2.8_4_8_5, -5.3_7_3_7], [-7.2_0_3_7, -4.5_5_0_5, -4.8_0_2_7], [-7.2_9_4_3, -4.2_6_1_1, -4.6_6_1_7]]   )
  __lowercase													=				torch.tensor([[0.4_9_8_7, 0.4_9_6_9, 0.9_9_9_9], [0.2_5_4_9, 0.5_4_9_8, 0.4_8_0_5], [0.5_4_9_8, 0.2_7_5_7, 0.0_5_6_9]]   )
 elif model_name == "deta-swin-large-o365":
  __lowercase													=				torch.tensor(
      [[-8.0_1_2_2, -3.5_7_2_0, -4.9_7_1_7], [-8.1_5_4_7, -3.6_8_8_6, -4.6_3_8_9], [-7.6_6_1_0, -3.6_1_9_4, -5.0_1_3_4]]   )
  __lowercase													=				torch.tensor([[0.2_5_2_3, 0.5_5_4_9, 0.4_8_8_1], [0.7_7_1_5, 0.4_1_4_9, 0.4_6_0_1], [0.5_5_0_3, 0.2_7_5_3, 0.0_5_7_5]]   )
 assert torch.allclose(outputs.logits[0, :3, :3]  ,       expected_logits.to(A__   )  ,       atol=1e-4   )
 assert torch.allclose(outputs.pred_boxes[0, :3, :3]  ,       expected_boxes.to(A__   )  ,       atol=1e-4   )
 print('''Everything ok!'''   )

 if pytorch_dump_folder_path:
  # Save model and processor
  logger.info(F"Saving PyTorch model and processor to {pytorch_dump_folder_path}..."   )
  Path(A__   ).mkdir(exist_ok=A__   )
  model.save_pretrained(A__   )
  processor.save_pretrained(A__   )

 # Push to hub
 if push_to_hub:
  print('''Pushing model and processor to hub...'''   )
  model.push_to_hub(F"jozhang97/{model_name}"   )
  processor.push_to_hub(F"jozhang97/{model_name}"   )


if __name__ == "__main__":
 lowerCAmelCase__      		=			argparse.ArgumentParser()

 parser.add_argument(
     '''--model_name''',
     type=str,
     default='''deta-swin-large''',
     choices=['''deta-swin-large''', '''deta-swin-large-o365'''],
     help='''Name of the model you\'d like to convert.''',
 )
 parser.add_argument(
     '''--pytorch_dump_folder_path''',
     default=None,
     type=str,
     help='''Path to the folder to output PyTorch model.''',
 )
 parser.add_argument(
     '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
 )
 lowerCAmelCase__      		=			parser.parse_args()
 convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


lowerCAmelCase__      		=			{'''configuration_vit_mae''': ['''VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ViTMAEConfig''']}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase__      		=			[
     '''VIT_MAE_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''ViTMAEForPreTraining''',
     '''ViTMAELayer''',
     '''ViTMAEModel''',
     '''ViTMAEPreTrainedModel''',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCAmelCase__      		=			[
     '''TFViTMAEForPreTraining''',
     '''TFViTMAEModel''',
     '''TFViTMAEPreTrainedModel''',
 ]

if TYPE_CHECKING:
 from .configuration_vit_mae import VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTMAEConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_vit_mae import (
      VIT_MAE_PRETRAINED_MODEL_ARCHIVE_LIST,
      ViTMAEForPreTraining,
      ViTMAELayer,
      ViTMAEModel,
      ViTMAEPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_vit_mae import TFViTMAEForPreTraining, TFViTMAEModel, TFViTMAEPreTrainedModel


else:
 import sys

 lowerCAmelCase__      		=			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

