




'''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()
_SCREAMING_SNAKE_CASE							:       Optional[int]						=			logging.get_logger(__name__)


def 					UpperCamelCase_(     snake_case							:	List[Any]					):

							'''simple docstring'''






							snake_case_      =				SwinConfig(
							    embed_dim=1_9_2							,    depths=(2, 2, 1_8, 2)							,    num_heads=(6, 1_2, 2_4, 4_8)							,    window_size=1_2							,    out_features=["stage2", "stage3", "stage4"]							,    )

							snake_case_      =				DetaConfig(
							    backbone_config=snake_case							,    num_queries=9_0_0							,    encoder_ffn_dim=2_0_4_8							,    decoder_ffn_dim=2_0_4_8							,    num_feature_levels=5							,    assign_first_stage=snake_case							,    with_box_refine=snake_case							,    two_stage=snake_case							,    )

							# set labels
							snake_case_      =				"huggingface/label-files"
							if "o365" in model_name:
														snake_case_      =				3_6_6
														snake_case_      =				"object365-id2label.json"
							else:
														snake_case_      =				9_1
														snake_case_      =				"coco-detection-id2label.json"

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

							return config


def 					UpperCamelCase_(     snake_case							:	Union[str, Any]					):

							'''simple docstring'''






							snake_case_      =				[]

							# 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 					UpperCamelCase_(     snake_case							:	Optional[Any]							,    snake_case							:	str							,    snake_case							:	Union[str, Any]					):

							'''simple docstring'''






							snake_case_      =				dct.pop(snake_case					)
							snake_case_      =				val


def 					UpperCamelCase_(     snake_case							:	Optional[Any]							,    snake_case							:	Optional[Any]					):

							'''simple docstring'''






							snake_case_      =				[int(backbone_config.embed_dim * 2**i					) for i in range(len(backbone_config.depths					)					)]
							for i in range(len(backbone_config.depths					)					):
														snake_case_      =				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)
																					snake_case_      =				state_dict.pop(f'backbone.0.body.layers.{i}.blocks.{j}.attn.qkv.weight'					)
																					snake_case_      =				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
																					snake_case_      =				in_proj_weight[:dim, :]
																					snake_case_      =				in_proj_bias[: dim]
																					snake_case_      =				in_proj_weight[
																					    dim : dim * 2, :
																					]
																					snake_case_      =				in_proj_bias[
																					    dim : dim * 2
																					]
																					snake_case_      =				in_proj_weight[
																					    -dim :, :
																					]
																					snake_case_      =				in_proj_bias[-dim :]
																					# fmt: on


def 					UpperCamelCase_(     snake_case							:	Tuple							,    snake_case							:	int					):

							'''simple docstring'''






							snake_case_      =				config.d_model
							for i in range(config.decoder_layers					):
														# read in weights + bias of input projection layer of self-attention
														snake_case_      =				state_dict.pop(f'transformer.decoder.layers.{i}.self_attn.in_proj_weight'					)
														snake_case_      =				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
														snake_case_      =				in_proj_weight[:hidden_size, :]
														snake_case_      =				in_proj_bias[:hidden_size]
														snake_case_      =				in_proj_weight[
														    hidden_size : hidden_size * 2, :
														]
														snake_case_      =				in_proj_bias[hidden_size : hidden_size * 2]
														snake_case_      =				in_proj_weight[-hidden_size:, :]
														snake_case_      =				in_proj_bias[-hidden_size:]


def 					UpperCamelCase_(     ):

							'''simple docstring'''






							snake_case_      =				"http://images.cocodataset.org/val2017/000000039769.jpg"
							snake_case_      =				Image.open(requests.get(snake_case							,    stream=snake_case					).raw					)

							return im


@torch.no_grad()
def 					UpperCamelCase_(     snake_case							:	Optional[int]							,    snake_case							:	int							,    snake_case							:	Optional[Any]					):

							'''simple docstring'''






							snake_case_      =				get_deta_config(snake_case					)

							# load original state dict
							if model_name == "deta-swin-large":
														snake_case_      =				hf_hub_download(repo_id="nielsr/deta-checkpoints"							,    filename="adet_swin_ft.pth"					)
							elif model_name == "deta-swin-large-o365":
														snake_case_      =				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'					)

							snake_case_      =				torch.load(snake_case							,    map_location="cpu"					)["model"]

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

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

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

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

							snake_case_      =				"cuda" if torch.cuda.is_available() else "cpu"
							model.to(snake_case					)

							# load image processor
							snake_case_      =				DetaImageProcessor(format="coco_detection"					)

							# verify our conversion on image
							snake_case_      =				prepare_img()
							snake_case_      =				processor(images=snake_case							,    return_tensors="pt"					)
							snake_case_      =				encoding["pixel_values"]
							snake_case_      =				model(pixel_values.to(snake_case					)					)

							# verify logits
							print("Logits:"							,    outputs.logits[0, :3, :3]					)
							print("Boxes:"							,    outputs.pred_boxes[0, :3, :3]					)
							if model_name == "deta-swin-large":
														snake_case_      =				torch.tensor(
														    [[-7.6308, -2.8485, -5.3737], [-7.2037, -4.5505, -4.8027], [-7.2943, -4.2611, -4.6617]]					)
														snake_case_      =				torch.tensor([[0.4987, 0.4969, 0.9999], [0.2549, 0.5498, 0.4805], [0.5498, 0.2757, 0.0569]]					)
							elif model_name == "deta-swin-large-o365":
														snake_case_      =				torch.tensor(
														    [[-8.0122, -3.5720, -4.9717], [-8.1547, -3.6886, -4.6389], [-7.6610, -3.6194, -5.0134]]					)
														snake_case_      =				torch.tensor([[0.2523, 0.5549, 0.4881], [0.7715, 0.4149, 0.4601], [0.5503, 0.2753, 0.0575]]					)
							assert torch.allclose(outputs.logits[0, :3, :3]							,    expected_logits.to(snake_case					)							,    atol=1e-4					)
							assert torch.allclose(outputs.pred_boxes[0, :3, :3]							,    expected_boxes.to(snake_case					)							,    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(snake_case					).mkdir(exist_ok=snake_case					)
														model.save_pretrained(snake_case					)
														processor.save_pretrained(snake_case					)

							# 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__":
		_SCREAMING_SNAKE_CASE							:       str						=			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."
		)
		_SCREAMING_SNAKE_CASE							:       Union[str, Any]						=			parser.parse_args()
		convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










'''simple docstring'''





import logging
import os

from .state import PartialState




class 			_snake_case		(     logging.LoggerAdapter							):




			@staticmethod
			def        lowerCAmelCase__    ( a__							)    ->    Optional[Any]:

										'''simple docstring'''




										snake_case_      =				PartialState()
										return not main_process_only or (main_process_only and state.is_main_process)




			def        lowerCAmelCase__    ( self    ,							a__    ,							a__    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										if PartialState._shared_state == {}:
																	raise RuntimeError(
																	    "You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility."							)
										snake_case_      =				kwargs.pop("main_process_only"    ,							a__							)
										snake_case_      =				kwargs.pop("in_order"    ,							a__							)

										if self.isEnabledFor(a__							):
																	if self._should_log(a__							):
																								snake_case_ ,				snake_case_      =				self.process(a__    ,							a__							)
																								self.logger.log(a__    ,							a__    ,							*a__    ,							**a__							)

																	elif in_order:
																								snake_case_      =				PartialState()
																								for i in range(state.num_processes							):
																															if i == state.process_index:
																																						snake_case_ ,				snake_case_      =				self.process(a__    ,							a__							)
																																						self.logger.log(a__    ,							a__    ,							*a__    ,							**a__							)
																															state.wait_for_everyone()


def 					UpperCamelCase_(     snake_case							:	str							,    snake_case							:	str = None					):

							'''simple docstring'''






							if log_level is None:
														snake_case_      =				os.environ.get("ACCELERATE_LOG_LEVEL"							,    snake_case					)
							snake_case_      =				logging.getLogger(snake_case					)
							if log_level is not None:
														logger.setLevel(log_level.upper()					)
														logger.root.setLevel(log_level.upper()					)
							return MultiProcessAdapter(snake_case							,    {}					)




