




"""simple docstring"""







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 DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor, ResNetConfig
from transformers.utils import logging


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





def      __snake_case (				_lowercase    ):


    """simple docstring"""







    if "resnet-50" in model_name:
        UpperCamelCase              =	ResNetConfig.from_pretrained('''microsoft/resnet-50'''    )
    elif "resnet-101" in model_name:
        UpperCamelCase              =	ResNetConfig.from_pretrained('''microsoft/resnet-101'''    )
    else:
        raise ValueError('''Model name should include either resnet50 or resnet101'''    )

    UpperCamelCase              =	DetrConfig(use_timm_backbone=_lowercase			,backbone_config=_lowercase    )

    # set label attributes
    UpperCamelCase              =	'''panoptic''' in model_name
    if is_panoptic:
        UpperCamelCase              =	250
    else:
        UpperCamelCase              =	91
        UpperCamelCase              =	'''huggingface/label-files'''
        UpperCamelCase              =	'''coco-detection-id2label.json'''
        UpperCamelCase              =	json.load(open(hf_hub_download(_lowercase			,_lowercase			,repo_type='''dataset'''    )			,'''r'''    )    )
        UpperCamelCase              =	{int(_lowercase    ): v for k, v in idalabel.items()}
        UpperCamelCase              =	idalabel
        UpperCamelCase              =	{v: k for k, v in idalabel.items()}

    return config, is_panoptic





def      __snake_case (				_lowercase    ):


    """simple docstring"""







    UpperCamelCase              =	[]

    # stem
    # fmt: off
    rename_keys.append(('''backbone.0.body.conv1.weight''', '''backbone.conv_encoder.model.embedder.embedder.convolution.weight''')    )
    rename_keys.append(('''backbone.0.body.bn1.weight''', '''backbone.conv_encoder.model.embedder.embedder.normalization.weight''')    )
    rename_keys.append(('''backbone.0.body.bn1.bias''', '''backbone.conv_encoder.model.embedder.embedder.normalization.bias''')    )
    rename_keys.append(('''backbone.0.body.bn1.running_mean''', '''backbone.conv_encoder.model.embedder.embedder.normalization.running_mean''')    )
    rename_keys.append(('''backbone.0.body.bn1.running_var''', '''backbone.conv_encoder.model.embedder.embedder.normalization.running_var''')    )
    # stages
    for stage_idx in range(len(config.backbone_config.depths    )    ):
        for layer_idx in range(config.backbone_config.depths[stage_idx]    ):
            # shortcut
            if layer_idx == 0:
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight',
                    )    )
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight',
                    )    )
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias',
                    )    )
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean',
                    )    )
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var',
                    )    )
            # 3 convs
            for i in range(3    ):
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight',
                    )    )
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight',
                    )    )
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias',
                    )    )
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean',
                    )    )
                rename_keys.append(
                    (
                        f'backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var',
                        f'backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var',
                    )    )
    # fmt: on

    for i in range(config.encoder_layers    ):
        # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
        rename_keys.append(
            (
                f'transformer.encoder.layers.{i}.self_attn.out_proj.weight',
                f'encoder.layers.{i}.self_attn.out_proj.weight',
            )    )
        rename_keys.append(
            (f'transformer.encoder.layers.{i}.self_attn.out_proj.bias', f'encoder.layers.{i}.self_attn.out_proj.bias')    )
        rename_keys.append((f'transformer.encoder.layers.{i}.linear1.weight', f'encoder.layers.{i}.fc1.weight')    )
        rename_keys.append((f'transformer.encoder.layers.{i}.linear1.bias', f'encoder.layers.{i}.fc1.bias')    )
        rename_keys.append((f'transformer.encoder.layers.{i}.linear2.weight', f'encoder.layers.{i}.fc2.weight')    )
        rename_keys.append((f'transformer.encoder.layers.{i}.linear2.bias', f'encoder.layers.{i}.fc2.bias')    )
        rename_keys.append(
            (f'transformer.encoder.layers.{i}.norm1.weight', f'encoder.layers.{i}.self_attn_layer_norm.weight')    )
        rename_keys.append(
            (f'transformer.encoder.layers.{i}.norm1.bias', f'encoder.layers.{i}.self_attn_layer_norm.bias')    )
        rename_keys.append(
            (f'transformer.encoder.layers.{i}.norm2.weight', f'encoder.layers.{i}.final_layer_norm.weight')    )
        rename_keys.append((f'transformer.encoder.layers.{i}.norm2.bias', f'encoder.layers.{i}.final_layer_norm.bias')    )
        # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
        rename_keys.append(
            (
                f'transformer.decoder.layers.{i}.self_attn.out_proj.weight',
                f'decoder.layers.{i}.self_attn.out_proj.weight',
            )    )
        rename_keys.append(
            (f'transformer.decoder.layers.{i}.self_attn.out_proj.bias', f'decoder.layers.{i}.self_attn.out_proj.bias')    )
        rename_keys.append(
            (
                f'transformer.decoder.layers.{i}.multihead_attn.out_proj.weight',
                f'decoder.layers.{i}.encoder_attn.out_proj.weight',
            )    )
        rename_keys.append(
            (
                f'transformer.decoder.layers.{i}.multihead_attn.out_proj.bias',
                f'decoder.layers.{i}.encoder_attn.out_proj.bias',
            )    )
        rename_keys.append((f'transformer.decoder.layers.{i}.linear1.weight', f'decoder.layers.{i}.fc1.weight')    )
        rename_keys.append((f'transformer.decoder.layers.{i}.linear1.bias', f'decoder.layers.{i}.fc1.bias')    )
        rename_keys.append((f'transformer.decoder.layers.{i}.linear2.weight', f'decoder.layers.{i}.fc2.weight')    )
        rename_keys.append((f'transformer.decoder.layers.{i}.linear2.bias', f'decoder.layers.{i}.fc2.bias')    )
        rename_keys.append(
            (f'transformer.decoder.layers.{i}.norm1.weight', f'decoder.layers.{i}.self_attn_layer_norm.weight')    )
        rename_keys.append(
            (f'transformer.decoder.layers.{i}.norm1.bias', f'decoder.layers.{i}.self_attn_layer_norm.bias')    )
        rename_keys.append(
            (f'transformer.decoder.layers.{i}.norm2.weight', f'decoder.layers.{i}.encoder_attn_layer_norm.weight')    )
        rename_keys.append(
            (f'transformer.decoder.layers.{i}.norm2.bias', f'decoder.layers.{i}.encoder_attn_layer_norm.bias')    )
        rename_keys.append(
            (f'transformer.decoder.layers.{i}.norm3.weight', f'decoder.layers.{i}.final_layer_norm.weight')    )
        rename_keys.append((f'transformer.decoder.layers.{i}.norm3.bias', f'decoder.layers.{i}.final_layer_norm.bias')    )

    # convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
    rename_keys.extend(
        [
            ('''input_proj.weight''', '''input_projection.weight'''),
            ('''input_proj.bias''', '''input_projection.bias'''),
            ('''query_embed.weight''', '''query_position_embeddings.weight'''),
            ('''transformer.decoder.norm.weight''', '''decoder.layernorm.weight'''),
            ('''transformer.decoder.norm.bias''', '''decoder.layernorm.bias'''),
            ('''class_embed.weight''', '''class_labels_classifier.weight'''),
            ('''class_embed.bias''', '''class_labels_classifier.bias'''),
            ('''bbox_embed.layers.0.weight''', '''bbox_predictor.layers.0.weight'''),
            ('''bbox_embed.layers.0.bias''', '''bbox_predictor.layers.0.bias'''),
            ('''bbox_embed.layers.1.weight''', '''bbox_predictor.layers.1.weight'''),
            ('''bbox_embed.layers.1.bias''', '''bbox_predictor.layers.1.bias'''),
            ('''bbox_embed.layers.2.weight''', '''bbox_predictor.layers.2.weight'''),
            ('''bbox_embed.layers.2.bias''', '''bbox_predictor.layers.2.bias'''),
        ]    )

    return rename_keys





def      __snake_case (				_lowercase			,_lowercase			,_lowercase    ):


    """simple docstring"""







    UpperCamelCase              =	state_dict.pop(_lowercase    )
    UpperCamelCase              =	val





def      __snake_case (				_lowercase			,_lowercase=False    ):


    """simple docstring"""







    UpperCamelCase              =	''''''
    if is_panoptic:
        UpperCamelCase              =	'''detr.'''

    # first: transformer encoder
    for i in range(6    ):
        # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
        UpperCamelCase              =	state_dict.pop(f'{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight'    )
        UpperCamelCase              =	state_dict.pop(f'{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias'    )
        # next, add query, keys and values (in that order) to the state dict
        UpperCamelCase              =	in_proj_weight[:256, :]
        UpperCamelCase              =	in_proj_bias[:256]
        UpperCamelCase              =	in_proj_weight[256:512, :]
        UpperCamelCase              =	in_proj_bias[256:512]
        UpperCamelCase              =	in_proj_weight[-256:, :]
        UpperCamelCase              =	in_proj_bias[-256:]
    # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
    for i in range(6    ):
        # read in weights + bias of input projection layer of self-attention
        UpperCamelCase              =	state_dict.pop(f'{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight'    )
        UpperCamelCase              =	state_dict.pop(f'{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias'    )
        # next, add query, keys and values (in that order) to the state dict
        UpperCamelCase              =	in_proj_weight[:256, :]
        UpperCamelCase              =	in_proj_bias[:256]
        UpperCamelCase              =	in_proj_weight[256:512, :]
        UpperCamelCase              =	in_proj_bias[256:512]
        UpperCamelCase              =	in_proj_weight[-256:, :]
        UpperCamelCase              =	in_proj_bias[-256:]
        # read in weights + bias of input projection layer of cross-attention
        UpperCamelCase              =	state_dict.pop(
            f'{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight'    )
        UpperCamelCase              =	state_dict.pop(f'{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias'    )
        # next, add query, keys and values (in that order) of cross-attention to the state dict
        UpperCamelCase              =	in_proj_weight_cross_attn[:256, :]
        UpperCamelCase              =	in_proj_bias_cross_attn[:256]
        UpperCamelCase              =	in_proj_weight_cross_attn[256:512, :]
        UpperCamelCase              =	in_proj_bias_cross_attn[256:512]
        UpperCamelCase              =	in_proj_weight_cross_attn[-256:, :]
        UpperCamelCase              =	in_proj_bias_cross_attn[-256:]





def      __snake_case (				):


    """simple docstring"""







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

    return im





@torch.no_grad()
def      __snake_case (				_lowercase			,_lowercase=None			,_lowercase=False    ):


    """simple docstring"""







    UpperCamelCase       ,	UpperCamelCase              =	get_detr_config(_lowercase    )

    # load original model from torch hub
    UpperCamelCase              =	{
        '''detr-resnet-50''': '''detr_resnet50''',
        '''detr-resnet-101''': '''detr_resnet101''',
    }
    logger.info(f'Converting model {model_name}...'    )
    UpperCamelCase              =	torch.hub.load('''facebookresearch/detr'''			,model_name_to_original_name[model_name]			,pretrained=_lowercase    ).eval()
    UpperCamelCase              =	detr.state_dict()
    # rename keys
    for src, dest in create_rename_keys(_lowercase    ):
        if is_panoptic:
            UpperCamelCase              =	'''detr.''' + src
        rename_key(_lowercase			,_lowercase			,_lowercase    )
    # query, key and value matrices need special treatment
    read_in_q_k_v(_lowercase			,is_panoptic=_lowercase    )
    # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
    UpperCamelCase              =	'''detr.model.''' if is_panoptic else '''model.'''
    for key in state_dict.copy().keys():
        if is_panoptic:
            if (
                key.startswith('''detr'''    )
                and not key.startswith('''class_labels_classifier'''    )
                and not key.startswith('''bbox_predictor'''    )
            ):
                UpperCamelCase              =	state_dict.pop(_lowercase    )
                UpperCamelCase              =	val
            elif "class_labels_classifier" in key or "bbox_predictor" in key:
                UpperCamelCase              =	state_dict.pop(_lowercase    )
                UpperCamelCase              =	val
            elif key.startswith('''bbox_attention'''    ) or key.startswith('''mask_head'''    ):
                continue
            else:
                UpperCamelCase              =	state_dict.pop(_lowercase    )
                UpperCamelCase              =	val
        else:
            if not key.startswith('''class_labels_classifier'''    ) and not key.startswith('''bbox_predictor'''    ):
                UpperCamelCase              =	state_dict.pop(_lowercase    )
                UpperCamelCase              =	val

    # finally, create HuggingFace model and load state dict
    UpperCamelCase              =	DetrForSegmentation(_lowercase    ) if is_panoptic else DetrForObjectDetection(_lowercase    )
    model.load_state_dict(_lowercase    )
    model.eval()

    # verify our conversion on an image
    UpperCamelCase              =	'''coco_panoptic''' if is_panoptic else '''coco_detection'''
    UpperCamelCase              =	DetrImageProcessor(format=_lowercase    )

    UpperCamelCase              =	processor(images=prepare_img()			,return_tensors='''pt'''    )
    UpperCamelCase              =	encoding['''pixel_values''']

    UpperCamelCase              =	detr(_lowercase    )
    UpperCamelCase              =	model(_lowercase    )

    assert torch.allclose(outputs.logits			,original_outputs['''pred_logits''']			,atol=1e-3    )
    assert torch.allclose(outputs.pred_boxes			,original_outputs['''pred_boxes''']			,atol=1e-3    )
    if is_panoptic:
        assert torch.allclose(outputs.pred_masks			,original_outputs['''pred_masks''']			,atol=1e-4    )
    print('''Looks ok!'''    )

    if pytorch_dump_folder_path is not None:
        # Save model and image processor
        logger.info(f'Saving PyTorch model and image processor to {pytorch_dump_folder_path}...'    )
        Path(_lowercase    ).mkdir(exist_ok=_lowercase    )
        model.save_pretrained(_lowercase    )
        processor.save_pretrained(_lowercase    )

    if push_to_hub:
        # Upload model and image processor to the hub
        logger.info('''Uploading PyTorch model and image processor to the hub...'''    )
        model.push_to_hub(f'nielsr/{model_name}'    )
        processor.push_to_hub(f'nielsr/{model_name}'    )


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

  parser.add_argument(
      '--model_name',
      default='detr-resnet-50',
      type=str,
      choices=['detr-resnet-50', 'detr-resnet-101'],
      help='Name of the DETR 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 to push the model to the hub or not.')
  SCREAMING_SNAKE_CASE_	     =       parser.parse_args()
  convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
'''simple docstring'''






import platform
from argparse import ArgumentParser

import huggingface_hub

from .. import __version__ as version
from ..utils import is_accelerate_available, is_torch_available, is_transformers_available, is_xformers_available
from . import BaseDiffusersCLICommand





def  _SCREAMING_SNAKE_CASE		(      __snake_case							:			Union[str, Any]				):
     return EnvironmentCommand()







class lowercase_       (		_UpperCamelCase			):







   """simple docstring"""

   @staticmethod
   def 	__UpperCAmelCase   (      UpperCamelCase__      :			ArgumentParser     )  ->			List[str]:
        _A    		=  parser.add_parser('env'     )
        download_parser.set_defaults(func=UpperCamelCase__     )
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			Optional[Any]:
        _A    		=  huggingface_hub.__version__

        _A    		=  'not installed'
        _A    		=  'NA'
        if is_torch_available():
             import torch

             _A    		=  torch.__version__
             _A    		=  torch.cuda.is_available()

        _A    		=  'not installed'
        if is_transformers_available():
             import transformers

             _A    		=  transformers.__version__

        _A    		=  'not installed'
        if is_accelerate_available():
             import accelerate

             _A    		=  accelerate.__version__

        _A    		=  'not installed'
        if is_xformers_available():
             import xformers

             _A    		=  xformers.__version__

        _A    		=  {
            '`diffusers` version': version,
            'Platform': platform.platform(),
            'Python version': platform.python_version(),
            'PyTorch version (GPU?)': f'{pt_version} ({pt_cuda_available})',
            'Huggingface_hub version': hub_version,
            'Transformers version': transformers_version,
            'Accelerate version': accelerate_version,
            'xFormers version': xformers_version,
            'Using GPU in script?': '<fill in>',
            'Using distributed or parallel set-up in script?': '<fill in>',
        }

        print('\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n'     )
        print(self.format_dict(UpperCamelCase__     )     )

        return info

   @staticmethod
   def 	__UpperCAmelCase   (      UpperCamelCase__      :			List[Any]     )  ->			Dict:
        return "\n".join([f'- {prop}: {val}' for prop, val in d.items()]     ) + "\n"



