



"""simple docstring"""





import argparse
import json
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import Callable, Dict, List, Tuple

import timm
import torch
import torch.nn as nn
from classy_vision.models.regnet import RegNet, RegNetParams, RegNetYaagf, RegNetYaagf, RegNetYaaagf
from huggingface_hub import cached_download, hf_hub_url
from torch import Tensor
from vissl.models.model_helpers import get_trunk_forward_outputs

from transformers import AutoImageProcessor, RegNetConfig, RegNetForImageClassification, RegNetModel
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase       	=				logging.get_logger()
@dataclass
class 				lowercase__       :







       '''simple docstring'''


       UpperCamelCase						=					42
       UpperCamelCase						=					field(default_factory=__snake_case				)
       UpperCamelCase						=					field(default_factory=__snake_case				)



       def 					lowercase__    (							self      :  str		,    _UpperCAmelCase      :  List[str]		,    _UpperCAmelCase      :  Tensor		,    _UpperCAmelCase      :  Tensor   )			->						Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_      			=							len(list(m.modules()   )   ) == 1 or isinstance(lowerCamelCase_		,    nn.Convad   ) or isinstance(lowerCamelCase_		,    nn.BatchNormad   )
             if has_not_submodules:
                   self.traced.append(lowerCamelCase_   )



       def __call__(							self      :  Tuple		,    _UpperCAmelCase      :  Tensor   )			->						Dict:


             '''simple docstring'''




             for m in self.module.modules():
                   self.handles.append(m.register_forward_hook(self._forward_hook   )   )
             self.module(lowerCamelCase_   )
             [x.remove() for x in self.handles]
             return self







       @property
       def 					lowercase__    (							self      :  List[str]   )			->						int:


             '''simple docstring'''




             return list(filter(lambda _UpperCAmelCase   : len(list(x.state_dict().keys()   )   ) > 0		,    self.traced   )   )



@dataclass
class 				lowercase__       :







       '''simple docstring'''


       UpperCamelCase						=					42
       UpperCamelCase						=					42
       UpperCamelCase						=					1
       UpperCamelCase						=					field(default_factory=__snake_case				)
       UpperCamelCase						=					field(default_factory=__snake_case				)
       UpperCamelCase						=					True



       def __call__(							self      :  Tuple		,    _UpperCAmelCase      :  Tensor   )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							Tracker(self.dest   )(lowerCamelCase_   ).parametrized
             UpperCAmelCase_      			=							Tracker(self.src   )(lowerCamelCase_   ).parametrized

             UpperCAmelCase_      			=							list(filter(lambda _UpperCAmelCase   : type(lowerCamelCase_   ) not in self.src_skip		,    lowerCamelCase_   )   )
             UpperCAmelCase_      			=							list(filter(lambda _UpperCAmelCase   : type(lowerCamelCase_   ) not in self.dest_skip		,    lowerCamelCase_   )   )

             if len(lowerCamelCase_   ) != len(lowerCamelCase_   ) and self.raise_if_mismatch:
                   raise Exception(
                       F"""Numbers of operations are different. Source module has {len(lowerCamelCase_   )} operations while"""
                       F""" destination module has {len(lowerCamelCase_   )}."""   )

             for dest_m, src_m in zip(lowerCamelCase_		,    lowerCamelCase_   ):
                   dest_m.load_state_dict(src_m.state_dict()   )
                   if self.verbose == 1:
                         print(F"""Transfered from={src_m} to={dest_m}"""   )



class 				lowercase__       (   nn.Module				):







       '''simple docstring'''


       def __init__(							self      :  List[str]		,    _UpperCAmelCase      :  nn.Module   )			->						List[str]:


             '''simple docstring'''




             super().__init__()

             UpperCAmelCase_      			=							[]
             # - get the stem
             feature_blocks.append(("conv1", model.stem)   )
             # - get all the feature blocks
             for k, v in model.trunk_output.named_children():
                   assert k.startswith("block"   ), F"""Unexpected layer name {k}"""
                   UpperCAmelCase_      			=							len(lowerCamelCase_   ) + 1
                   feature_blocks.append((F"""res{block_index}""", v)   )

             UpperCAmelCase_      			=							nn.ModuleDict(lowerCamelCase_   )







       def 					lowercase__    (							self      :  Union[str, Any]		,    _UpperCAmelCase      :  Tensor   )			->						List[str]:


             '''simple docstring'''




             return get_trunk_forward_outputs(
                 lowerCamelCase_		,    out_feat_keys=lowerCamelCase_		,    feature_blocks=self._feature_blocks		,    )



class 				lowercase__       (   __snake_case				):







       '''simple docstring'''


       def 					lowercase__    (							self      :  Dict		,    _UpperCAmelCase      :  str   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							x.split("-"   )
             return x_split[0] + x_split[1] + "_" + "".join(x_split[2:]   )







       def __getitem__(							self      :  Union[str, Any]		,    _UpperCAmelCase      :  str   )			->						Callable[[], Tuple[nn.Module, Dict]]:


             '''simple docstring'''




             if x not in self:
                   UpperCAmelCase_      			=							self.convert_name_to_timm(lowerCamelCase_   )
                   UpperCAmelCase_      			=							partial(lambda: (timm.create_model(lowerCamelCase_		,    pretrained=lowerCamelCase_   ).eval(), None)   )

             else:
                   UpperCAmelCase_      			=							super().__getitem__(lowerCamelCase_   )

             return val



class 				lowercase__       (   __snake_case				):







       '''simple docstring'''


       def __getitem__(							self      :  Union[str, Any]		,    _UpperCAmelCase      :  str   )			->						Callable[[], nn.Module]:


             '''simple docstring'''




             if "seer" in x and "in1k" not in x:
                   UpperCAmelCase_      			=							RegNetModel
             else:
                   UpperCAmelCase_      			=							RegNetForImageClassification
             return val



def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__				):




      for from_key, to_key in keys:
            UpperCAmelCase_      			=							from_state_dict[from_key].clone()
            print(f"""Copied key={from_key} to={to_key}"""				)
      return to_state_dict
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__ = True     ,      ):
      print(f"""Converting {name}..."""				)
      with torch.no_grad():
            UpperCAmelCase_      			=							from_model_func()
            UpperCAmelCase_      			=							our_model_func(_a				).eval()
            UpperCAmelCase_      			=							ModuleTransfer(src=_a     ,      dest=_a     ,      raise_if_mismatch=_a				)
            UpperCAmelCase_      			=							torch.randn((1, 3, 224, 224)				)
            module_transfer(_a				)

      if from_state_dict is not None:
            UpperCAmelCase_      			=							[]
            # for seer - in1k finetuned we have to manually copy the head
            if "seer" in name and "in1k" in name:
                  UpperCAmelCase_      			=							[("""0.clf.0.weight""", """classifier.1.weight"""), ("""0.clf.0.bias""", """classifier.1.bias""")]
            UpperCAmelCase_      			=							manually_copy_vissl_head(_a     ,      our_model.state_dict()     ,      _a				)
            our_model.load_state_dict(_a				)

      UpperCAmelCase_      			=							our_model(_a     ,      output_hidden_states=_a				)
      UpperCAmelCase_      			=							(
          our_outputs.logits if isinstance(_a     ,      _a				) else our_outputs.last_hidden_state
      )

      UpperCAmelCase_      			=							from_model(_a				)
      UpperCAmelCase_      			=							from_output[-1] if type(_a				) is list else from_output

      # now since I don't want to use any config files, vissl seer model doesn't actually have an head, so let's just check the last hidden state
      if "seer" in name and "in1k" in name:
            UpperCAmelCase_      			=							our_outputs.hidden_states[-1]

      assert torch.allclose(_a     ,      _a				), "The model logits don't match the original one."

      if push_to_hub:
            our_model.push_to_hub(
                repo_path_or_name=save_directory / name     ,      commit_message="Add model"     ,      use_temp_dir=_a     ,      )

            UpperCAmelCase_      			=							224 if """seer""" not in name else 384
            # we can use the convnext one
            UpperCAmelCase_      			=							AutoImageProcessor.from_pretrained("facebook/convnext-base-224-22k-1k"     ,      size=_a				)
            image_processor.push_to_hub(
                repo_path_or_name=save_directory / name     ,      commit_message="Add image processor"     ,      use_temp_dir=_a     ,      )

            print(f"""Pushed {name}"""				)
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__ = None     ,      lowerCAmelCase__ = True				):
      UpperCAmelCase_      			=							"""imagenet-1k-id2label.json"""
      UpperCAmelCase_      			=							1000
      UpperCAmelCase_      			=							(1, num_labels)

      UpperCAmelCase_      			=							"""huggingface/label-files"""
      UpperCAmelCase_      			=							num_labels
      UpperCAmelCase_      			=							json.load(open(cached_download(hf_hub_url(_a     ,      _a     ,      repo_type="dataset"				)				)     ,      "r"				)				)
      UpperCAmelCase_      			=							{int(_a				): v for k, v in idalabel.items()}

      UpperCAmelCase_      			=							idalabel
      UpperCAmelCase_      			=							{v: k for k, v in idalabel.items()}

      UpperCAmelCase_      			=							partial(_a     ,      num_labels=_a     ,      idalabel=_a     ,      labelaid=_a				)

      UpperCAmelCase_      			=							{
          """regnet-x-002""": ImageNetPreTrainedConfig(
              depths=[1, 1, 4, 7]     ,      hidden_sizes=[24, 56, 152, 368]     ,      groups_width=8     ,      layer_type="x"				),
          """regnet-x-004""": ImageNetPreTrainedConfig(
              depths=[1, 2, 7, 12]     ,      hidden_sizes=[32, 64, 160, 384]     ,      groups_width=16     ,      layer_type="x"				),
          """regnet-x-006""": ImageNetPreTrainedConfig(
              depths=[1, 3, 5, 7]     ,      hidden_sizes=[48, 96, 240, 528]     ,      groups_width=24     ,      layer_type="x"				),
          """regnet-x-008""": ImageNetPreTrainedConfig(
              depths=[1, 3, 7, 5]     ,      hidden_sizes=[64, 128, 288, 672]     ,      groups_width=16     ,      layer_type="x"				),
          """regnet-x-016""": ImageNetPreTrainedConfig(
              depths=[2, 4, 10, 2]     ,      hidden_sizes=[72, 168, 408, 912]     ,      groups_width=24     ,      layer_type="x"				),
          """regnet-x-032""": ImageNetPreTrainedConfig(
              depths=[2, 6, 15, 2]     ,      hidden_sizes=[96, 192, 432, 1008]     ,      groups_width=48     ,      layer_type="x"				),
          """regnet-x-040""": ImageNetPreTrainedConfig(
              depths=[2, 5, 14, 2]     ,      hidden_sizes=[80, 240, 560, 1360]     ,      groups_width=40     ,      layer_type="x"				),
          """regnet-x-064""": ImageNetPreTrainedConfig(
              depths=[2, 4, 10, 1]     ,      hidden_sizes=[168, 392, 784, 1624]     ,      groups_width=56     ,      layer_type="x"				),
          """regnet-x-080""": ImageNetPreTrainedConfig(
              depths=[2, 5, 15, 1]     ,      hidden_sizes=[80, 240, 720, 1920]     ,      groups_width=120     ,      layer_type="x"				),
          """regnet-x-120""": ImageNetPreTrainedConfig(
              depths=[2, 5, 11, 1]     ,      hidden_sizes=[224, 448, 896, 2240]     ,      groups_width=112     ,      layer_type="x"				),
          """regnet-x-160""": ImageNetPreTrainedConfig(
              depths=[2, 6, 13, 1]     ,      hidden_sizes=[256, 512, 896, 2048]     ,      groups_width=128     ,      layer_type="x"				),
          """regnet-x-320""": ImageNetPreTrainedConfig(
              depths=[2, 7, 13, 1]     ,      hidden_sizes=[336, 672, 1344, 2520]     ,      groups_width=168     ,      layer_type="x"				),
          # y variant
          """regnet-y-002""": ImageNetPreTrainedConfig(depths=[1, 1, 4, 7]     ,      hidden_sizes=[24, 56, 152, 368]     ,      groups_width=8				),
          """regnet-y-004""": ImageNetPreTrainedConfig(
              depths=[1, 3, 6, 6]     ,      hidden_sizes=[48, 104, 208, 440]     ,      groups_width=8				),
          """regnet-y-006""": ImageNetPreTrainedConfig(
              depths=[1, 3, 7, 4]     ,      hidden_sizes=[48, 112, 256, 608]     ,      groups_width=16				),
          """regnet-y-008""": ImageNetPreTrainedConfig(
              depths=[1, 3, 8, 2]     ,      hidden_sizes=[64, 128, 320, 768]     ,      groups_width=16				),
          """regnet-y-016""": ImageNetPreTrainedConfig(
              depths=[2, 6, 17, 2]     ,      hidden_sizes=[48, 120, 336, 888]     ,      groups_width=24				),
          """regnet-y-032""": ImageNetPreTrainedConfig(
              depths=[2, 5, 13, 1]     ,      hidden_sizes=[72, 216, 576, 1512]     ,      groups_width=24				),
          """regnet-y-040""": ImageNetPreTrainedConfig(
              depths=[2, 6, 12, 2]     ,      hidden_sizes=[128, 192, 512, 1088]     ,      groups_width=64				),
          """regnet-y-064""": ImageNetPreTrainedConfig(
              depths=[2, 7, 14, 2]     ,      hidden_sizes=[144, 288, 576, 1296]     ,      groups_width=72				),
          """regnet-y-080""": ImageNetPreTrainedConfig(
              depths=[2, 4, 10, 1]     ,      hidden_sizes=[168, 448, 896, 2016]     ,      groups_width=56				),
          """regnet-y-120""": ImageNetPreTrainedConfig(
              depths=[2, 5, 11, 1]     ,      hidden_sizes=[224, 448, 896, 2240]     ,      groups_width=112				),
          """regnet-y-160""": ImageNetPreTrainedConfig(
              depths=[2, 4, 11, 1]     ,      hidden_sizes=[224, 448, 1232, 3024]     ,      groups_width=112				),
          """regnet-y-320""": ImageNetPreTrainedConfig(
              depths=[2, 5, 12, 1]     ,      hidden_sizes=[232, 696, 1392, 3712]     ,      groups_width=232				),
          # models created by SEER -> https://arxiv.org/abs/2202.08360
          """regnet-y-320-seer""": RegNetConfig(depths=[2, 5, 12, 1]     ,      hidden_sizes=[232, 696, 1392, 3712]     ,      groups_width=232				),
          """regnet-y-640-seer""": RegNetConfig(depths=[2, 5, 12, 1]     ,      hidden_sizes=[328, 984, 1968, 4920]     ,      groups_width=328				),
          """regnet-y-1280-seer""": RegNetConfig(
              depths=[2, 7, 17, 1]     ,      hidden_sizes=[528, 1056, 2904, 7392]     ,      groups_width=264				),
          """regnet-y-2560-seer""": RegNetConfig(
              depths=[3, 7, 16, 1]     ,      hidden_sizes=[640, 1696, 2544, 5088]     ,      groups_width=640				),
          """regnet-y-10b-seer""": ImageNetPreTrainedConfig(
              depths=[2, 7, 17, 1]     ,      hidden_sizes=[2020, 4040, 11110, 28280]     ,      groups_width=1010				),
          # finetuned on imagenet
          """regnet-y-320-seer-in1k""": ImageNetPreTrainedConfig(
              depths=[2, 5, 12, 1]     ,      hidden_sizes=[232, 696, 1392, 3712]     ,      groups_width=232				),
          """regnet-y-640-seer-in1k""": ImageNetPreTrainedConfig(
              depths=[2, 5, 12, 1]     ,      hidden_sizes=[328, 984, 1968, 4920]     ,      groups_width=328				),
          """regnet-y-1280-seer-in1k""": ImageNetPreTrainedConfig(
              depths=[2, 7, 17, 1]     ,      hidden_sizes=[528, 1056, 2904, 7392]     ,      groups_width=264				),
          """regnet-y-2560-seer-in1k""": ImageNetPreTrainedConfig(
              depths=[3, 7, 16, 1]     ,      hidden_sizes=[640, 1696, 2544, 5088]     ,      groups_width=640				),
          """regnet-y-10b-seer-in1k""": ImageNetPreTrainedConfig(
              depths=[2, 7, 17, 1]     ,      hidden_sizes=[2020, 4040, 11110, 28280]     ,      groups_width=1010				),
      }

      UpperCAmelCase_      			=							NameToOurModelFuncMap()
      UpperCAmelCase_      			=							NameToFromModelFuncMap()
      # add seer weights logic

      def load_using_classy_vision(lowerCAmelCase__     ,      lowerCAmelCase__				) -> Tuple[nn.Module, Dict]:
            UpperCAmelCase_      			=							torch.hub.load_state_dict_from_url(_a     ,      model_dir=str(_a				)     ,      map_location="cpu"				)
            UpperCAmelCase_      			=							model_func()
            # check if we have a head, if yes add it
            UpperCAmelCase_      			=							files["""classy_state_dict"""]["""base_model"""]["""model"""]
            UpperCAmelCase_      			=							model_state_dict["""trunk"""]
            model.load_state_dict(_a				)
            return model.eval(), model_state_dict["heads"]

      # pretrained
      UpperCAmelCase_      			=							partial(
          _a     ,      "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet32d/seer_regnet32gf_model_iteration244000.torch"     ,      lambda: FakeRegNetVisslWrapper(RegNetYaagf()				)     ,      )

      UpperCAmelCase_      			=							partial(
          _a     ,      "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet64/seer_regnet64gf_model_final_checkpoint_phase0.torch"     ,      lambda: FakeRegNetVisslWrapper(RegNetYaagf()				)     ,      )

      UpperCAmelCase_      			=							partial(
          _a     ,      "https://dl.fbaipublicfiles.com/vissl/model_zoo/swav_ig1b_regnet128Gf_cnstant_bs32_node16_sinkhorn10_proto16k_syncBN64_warmup8k/model_final_checkpoint_phase0.torch"     ,      lambda: FakeRegNetVisslWrapper(RegNetYaaagf()				)     ,      )

      UpperCAmelCase_      			=							partial(
          _a     ,      "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet10B/model_iteration124500_conso.torch"     ,      lambda: FakeRegNetVisslWrapper(
              RegNet(RegNetParams(depth=27     ,      group_width=1010     ,      w_a=1744     ,      w_a=620.83     ,      w_m=2.52				)				)				)     ,      )

      # IN1K finetuned
      UpperCAmelCase_      			=							partial(
          _a     ,      "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet32_finetuned_in1k_model_final_checkpoint_phase78.torch"     ,      lambda: FakeRegNetVisslWrapper(RegNetYaagf()				)     ,      )

      UpperCAmelCase_      			=							partial(
          _a     ,      "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet64_finetuned_in1k_model_final_checkpoint_phase78.torch"     ,      lambda: FakeRegNetVisslWrapper(RegNetYaagf()				)     ,      )

      UpperCAmelCase_      			=							partial(
          _a     ,      "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet128_finetuned_in1k_model_final_checkpoint_phase78.torch"     ,      lambda: FakeRegNetVisslWrapper(RegNetYaaagf()				)     ,      )

      UpperCAmelCase_      			=							partial(
          _a     ,      "https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_10b_finetuned_in1k_model_phase28_conso.torch"     ,      lambda: FakeRegNetVisslWrapper(
              RegNet(RegNetParams(depth=27     ,      group_width=1010     ,      w_a=1744     ,      w_a=620.83     ,      w_m=2.52				)				)				)     ,      )

      if model_name:
            convert_weight_and_push(
                _a     ,      names_to_from_model_map[model_name]     ,      names_to_ours_model_map[model_name]     ,      names_to_config[model_name]     ,      _a     ,      _a     ,      )
      else:
            for model_name, config in names_to_config.items():
                  convert_weight_and_push(
                      _a     ,      names_to_from_model_map[model_name]     ,      names_to_ours_model_map[model_name]     ,      _a     ,      _a     ,      _a     ,      )
      return config, expected_shape


if __name__ == "__main__":
  lowerCamelCase       	=				argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--model_name""",
      default=None,
      type=str,
      help=(
          """The name of the model you wish to convert, it must be one of the supported regnet* architecture,"""
          """ currently: regnetx-*, regnety-*. If `None`, all of them will the converted."""
      ),
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""",
      default=None,
      type=Path,
      required=True,
      help="""Path to the output PyTorch model directory.""",
  )
  parser.add_argument(
      """--push_to_hub""",
      default=True,
      type=bool,
      required=False,
      help="""If True, push model and image processor to the hub.""",
  )

  lowerCamelCase       	=				parser.parse_args()

  lowerCamelCase       	=				args.pytorch_dump_folder_path
  pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
  convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)









"""simple docstring"""





import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

  if is_vision_available():
    from transformers import MaskaFormerImageProcessor

if is_vision_available():
  from PIL import Image
class 				lowercase__       :







       '''simple docstring'''


       def __init__(							self      :  Any		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  Union[str, Any]=2		,    _UpperCAmelCase      :  Any=True		,    _UpperCAmelCase      :  int=False		,    _UpperCAmelCase      :  str=10		,    _UpperCAmelCase      :  int=3		,    _UpperCAmelCase      :  Optional[int]=32 * 8		,    _UpperCAmelCase      :  str=32 * 8		,    _UpperCAmelCase      :  Any=4		,    _UpperCAmelCase      :  List[Any]=64		,    )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							parent
             UpperCAmelCase_      			=							batch_size
             UpperCAmelCase_      			=							is_training
             UpperCAmelCase_      			=							use_auxiliary_loss
             UpperCAmelCase_      			=							num_queries
             UpperCAmelCase_      			=							num_channels
             UpperCAmelCase_      			=							min_size
             UpperCAmelCase_      			=							max_size
             UpperCAmelCase_      			=							num_labels
             UpperCAmelCase_      			=							hidden_dim
             UpperCAmelCase_      			=							hidden_dim



       def 					lowercase__    (							self      :  int   )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]   ).to(
                 _UpperCAmelCase   )

             UpperCAmelCase_      			=							torch.ones([self.batch_size, self.min_size, self.max_size]		,    device=_UpperCAmelCase   )

             UpperCAmelCase_      			=							(
                 torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]		,    device=_UpperCAmelCase   ) > 0.5
             ).float()
             UpperCAmelCase_      			=							(torch.rand((self.batch_size, self.num_labels)		,    device=_UpperCAmelCase   ) > 0.5).long()

             UpperCAmelCase_      			=							self.get_config()
             return config, pixel_values, pixel_mask, mask_labels, class_labels



       def 					lowercase__    (							self      :  Any   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerConfig(
                 hidden_size=self.hidden_dim		,    )
             UpperCAmelCase_      			=							self.num_queries
             UpperCAmelCase_      			=							self.num_labels

             UpperCAmelCase_      			=							[1, 1, 1, 1]
             UpperCAmelCase_      			=							self.num_channels

             UpperCAmelCase_      			=							64
             UpperCAmelCase_      			=							128
             UpperCAmelCase_      			=							self.hidden_dim
             UpperCAmelCase_      			=							self.hidden_dim
             UpperCAmelCase_      			=							self.hidden_dim
             return config



       def 					lowercase__    (							self      :  Dict   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self.prepare_config_and_inputs()
             UpperCAmelCase_      			=							{"pixel_values": pixel_values, "pixel_mask": pixel_mask}
             return config, inputs_dict



       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  Union[str, Any]   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_      			=							output.encoder_hidden_states
             UpperCAmelCase_      			=							output.pixel_decoder_hidden_states
             UpperCAmelCase_      			=							output.transformer_decoder_hidden_states

             self.parent.assertTrue(len(_UpperCAmelCase   )		,    len(config.backbone_config.depths   )   )
             self.parent.assertTrue(len(_UpperCAmelCase   )		,    len(config.backbone_config.depths   )   )
             self.parent.assertTrue(len(_UpperCAmelCase   )		,    config.decoder_layers   )



       def 					lowercase__    (							self      :  Tuple		,    _UpperCAmelCase      :  List[Any]		,    _UpperCAmelCase      :  Any		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  Optional[int]=False   )			->						str:


             '''simple docstring'''




             with torch.no_grad():
                   UpperCAmelCase_      			=							MaskaFormerModel(config=_UpperCAmelCase   )
                   model.to(_UpperCAmelCase   )
                   model.eval()

                   UpperCAmelCase_      			=							model(pixel_values=_UpperCAmelCase		,    pixel_mask=_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(_UpperCAmelCase		,    output_hidden_states=_UpperCAmelCase   )

             self.parent.assertEqual(
                 output.transformer_decoder_last_hidden_state.shape		,    (self.batch_size, self.num_queries, self.hidden_dim)		,    )
             # let's ensure the other two hidden state exists
             self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None   )
             self.parent.assertTrue(output.encoder_last_hidden_state is not None   )

             if output_hidden_states:
                   self.check_output_hidden_state(_UpperCAmelCase		,    _UpperCAmelCase   )







       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  Optional[int]		,    _UpperCAmelCase      :  Dict		,    _UpperCAmelCase      :  Optional[int]		,    _UpperCAmelCase      :  Union[str, Any]		,    _UpperCAmelCase      :  int   )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation(config=_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.eval()

             def comm_check_on_output(_UpperCAmelCase      :  List[Any]   ):
                   # let's still check that all the required stuff is there
                   self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None   )
                   self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None   )
                   self.parent.assertTrue(result.encoder_last_hidden_state is not None   )
                   # okay, now we need to check the logits shape
                   # due to the encoder compression, masks have a //4 spatial size
                   self.parent.assertEqual(
                       result.masks_queries_logits.shape		,    (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)		,    )
                   # + 1 for null class
                   self.parent.assertEqual(
                       result.class_queries_logits.shape		,    (self.batch_size, self.num_queries, self.num_labels + 1)   )

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(pixel_values=_UpperCAmelCase		,    pixel_mask=_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(_UpperCAmelCase   )

                   comm_check_on_output(_UpperCAmelCase   )

                   UpperCAmelCase_      			=							model(
                       pixel_values=_UpperCAmelCase		,    pixel_mask=_UpperCAmelCase		,    mask_labels=_UpperCAmelCase		,    class_labels=_UpperCAmelCase   )

             comm_check_on_output(_UpperCAmelCase   )

             self.parent.assertTrue(result.loss is not None   )
             self.parent.assertEqual(result.loss.shape		,    torch.Size([1]   )   )



@require_torch
class 				lowercase__       (   SCREAMING_SNAKE_CASE    , SCREAMING_SNAKE_CASE    , unittest.TestCase				):







       '''simple docstring'''


       UpperCamelCase						=					(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
       UpperCamelCase						=					{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

       UpperCamelCase						=					False
       UpperCamelCase						=					False
       UpperCamelCase						=					False
       UpperCamelCase						=					False



       def 					lowercase__    (							self      :  Tuple   )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerModelTester(self   )
             UpperCAmelCase_      			=							ConfigTester(self		,    config_class=_UpperCAmelCase		,    has_text_modality=_UpperCAmelCase   )



       def 					lowercase__    (							self      :  List[str]   )			->						Dict:


             '''simple docstring'''




             self.config_tester.run_common_tests()



       def 					lowercase__    (							self      :  Dict   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()
             self.model_tester.create_and_check_maskaformer_model(_UpperCAmelCase		,    **_UpperCAmelCase		,    output_hidden_states=_UpperCAmelCase   )



       def 					lowercase__    (							self      :  Optional[int]   )			->						Dict:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*_UpperCAmelCase   )



       @unittest.skip(reason="Mask2Former does not use inputs_embeds"   )
       def 					lowercase__    (							self      :  Tuple   )			->						Any:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Mask2Former does not have a get_input_embeddings method"   )
       def 					lowercase__    (							self      :  Tuple   )			->						Optional[int]:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Mask2Former is not a generative model"   )
       def 					lowercase__    (							self      :  str   )			->						Dict:


             '''simple docstring'''




             pass



       @unittest.skip(reason="Mask2Former does not use token embeddings"   )
       def 					lowercase__    (							self      :  Union[str, Any]   )			->						Union[str, Any]:


             '''simple docstring'''




             pass



       @require_torch_multi_gpu
       @unittest.skip(
           reason="Mask2Former has some layers using `add_module` which doesn't work well with `nn.DataParallel`"   )
       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''simple docstring'''




             pass



       @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."   )
       def 					lowercase__    (							self      :  List[str]   )			->						List[str]:


             '''simple docstring'''




             pass



       def 					lowercase__    (							self      :  Tuple   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   UpperCAmelCase_      			=							model_class(_UpperCAmelCase   )
                   UpperCAmelCase_      			=							inspect.signature(model.forward   )
                   # signature.parameters is an OrderedDict => so arg_names order is deterministic
                   UpperCAmelCase_      			=							[*signature.parameters.keys()]

                   UpperCAmelCase_      			=							["pixel_values"]
                   self.assertListEqual(arg_names[:1]		,    _UpperCAmelCase   )



       @slow
       def 					lowercase__    (							self      :  List[Any]   )			->						Dict:


             '''simple docstring'''




             for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
                   UpperCAmelCase_      			=							MaskaFormerModel.from_pretrained(_UpperCAmelCase   )
                   self.assertIsNotNone(_UpperCAmelCase   )



       def 					lowercase__    (							self      :  List[str]   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_      			=							(self.model_tester.min_size,) * 2
             UpperCAmelCase_      			=							{
                 "pixel_values": torch.randn((2, 3, *size)		,    device=_UpperCAmelCase   ),
                 "mask_labels": torch.randn((2, 10, *size)		,    device=_UpperCAmelCase   ),
                 "class_labels": torch.zeros(2		,    10		,    device=_UpperCAmelCase   ).long(),
             }
             UpperCAmelCase_      			=							self.model_tester.get_config()

             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation(_UpperCAmelCase   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							model(**_UpperCAmelCase   )
             self.assertTrue(outputs.loss is not None   )



       def 					lowercase__    (							self      :  Union[str, Any]   )			->						str:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()
             self.model_tester.create_and_check_maskaformer_model(_UpperCAmelCase		,    **_UpperCAmelCase		,    output_hidden_states=_UpperCAmelCase   )



       def 					lowercase__    (							self      :  Union[str, Any]   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   UpperCAmelCase_      			=							model_class(_UpperCAmelCase   ).to(_UpperCAmelCase   )
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase		,    output_attentions=_UpperCAmelCase   )
                   self.assertTrue(outputs.attentions is not None   )



       def 					lowercase__    (							self      :  Union[str, Any]   )			->						List[str]:


             '''simple docstring'''




             if not self.model_tester.is_training:
                   return

             UpperCAmelCase_      			=							self.all_model_classes[1]
             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()

             UpperCAmelCase_      			=							model_class(_UpperCAmelCase   )
             model.to(_UpperCAmelCase   )
             model.train()

             UpperCAmelCase_      			=							model(_UpperCAmelCase		,    mask_labels=_UpperCAmelCase		,    class_labels=_UpperCAmelCase   ).loss
             loss.backward()







       def 					lowercase__    (							self      :  str   )			->						Optional[int]:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.all_model_classes[1]
             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							self.model_tester.prepare_config_and_inputs()
             UpperCAmelCase_      			=							True
             UpperCAmelCase_      			=							True

             UpperCAmelCase_      			=							model_class(_UpperCAmelCase   ).to(_UpperCAmelCase   )
             model.train()

             UpperCAmelCase_      			=							model(_UpperCAmelCase		,    mask_labels=_UpperCAmelCase		,    class_labels=_UpperCAmelCase   )

             UpperCAmelCase_      			=							outputs.encoder_hidden_states[0]
             encoder_hidden_states.retain_grad()

             UpperCAmelCase_      			=							outputs.pixel_decoder_hidden_states[0]
             pixel_decoder_hidden_states.retain_grad()

             UpperCAmelCase_      			=							outputs.transformer_decoder_hidden_states[0]
             transformer_decoder_hidden_states.retain_grad()

             UpperCAmelCase_      			=							outputs.attentions[0]
             attentions.retain_grad()

             outputs.loss.backward(retain_graph=_UpperCAmelCase   )

             self.assertIsNotNone(encoder_hidden_states.grad   )
             self.assertIsNotNone(pixel_decoder_hidden_states.grad   )
             self.assertIsNotNone(transformer_decoder_hidden_states.grad   )
             self.assertIsNotNone(attentions.grad   )



lowerCamelCase       	=						1e-4
def        a__     (				):
      UpperCAmelCase_      			=							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"				)
      return image


@require_vision
@slow
class 				lowercase__       (   unittest.TestCase				):







       '''simple docstring'''


       @cached_property
       def 					lowercase__    (							self      :  Optional[Any]   )			->						int:


             '''simple docstring'''




             return "facebook/mask2former-swin-small-coco-instance"



       @cached_property
       def 					lowercase__    (							self      :  Dict   )			->						Optional[int]:


             '''simple docstring'''




             return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints   ) if is_vision_available() else None



       def 					lowercase__    (							self      :  Optional[Any]   )			->						Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerModel.from_pretrained(self.model_checkpoints   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							self.default_image_processor
             UpperCAmelCase_      			=							prepare_img()
             UpperCAmelCase_      			=							image_processor(_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							inputs["pixel_values"].shape
             # check size is divisible by 32
             self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
             # check size
             self.assertEqual(_UpperCAmelCase		,    (1, 3, 384, 384)   )

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]   ).to(_UpperCAmelCase   )
             self.assertTrue(
                 torch.allclose(
                     outputs.encoder_last_hidden_state[0, 0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]   ).to(_UpperCAmelCase   )
             self.assertTrue(
                 torch.allclose(
                     outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )

             UpperCAmelCase_      			=							torch.tensor(
                 [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]   ).to(_UpperCAmelCase   )
             self.assertTrue(
                 torch.allclose(
                     outputs.transformer_decoder_last_hidden_state[0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )



       def 					lowercase__    (							self      :  int   )			->						List[str]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_UpperCAmelCase   ).eval()
             UpperCAmelCase_      			=							self.default_image_processor
             UpperCAmelCase_      			=							prepare_img()
             UpperCAmelCase_      			=							image_processor(_UpperCAmelCase		,    return_tensors="pt"   ).to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							inputs["pixel_values"].shape
             # check size is divisible by 32
             self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
             # check size
             self.assertEqual(_UpperCAmelCase		,    (1, 3, 384, 384)   )

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )
             # masks_queries_logits
             UpperCAmelCase_      			=							outputs.masks_queries_logits
             self.assertEqual(
                 masks_queries_logits.shape		,    (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)   )
             UpperCAmelCase_      			=							[
                 [-8.7839, -9.0056, -8.8121],
                 [-7.4104, -7.0313, -6.5401],
                 [-6.6105, -6.3427, -6.4675],
             ]
             UpperCAmelCase_      			=							torch.tensor(_UpperCAmelCase   ).to(_UpperCAmelCase   )
             self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )
             # class_queries_logits
             UpperCAmelCase_      			=							outputs.class_queries_logits
             self.assertEqual(class_queries_logits.shape		,    (1, model.config.num_queries, model.config.num_labels + 1)   )
             UpperCAmelCase_      			=							torch.tensor(
                 [
                     [1.8324, -8.0835, -4.1922],
                     [0.8450, -9.0050, -3.6053],
                     [0.3045, -7.7293, -3.0275],
                 ]   ).to(_UpperCAmelCase   )
             self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]		,    _UpperCAmelCase		,    atol=_UpperCAmelCase   )   )







       def 					lowercase__    (							self      :  List[Any]   )			->						Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_UpperCAmelCase   ).eval()
             UpperCAmelCase_      			=							self.default_image_processor

             UpperCAmelCase_      			=							image_processor(
                 [np.zeros((3, 800, 1333)   ), np.zeros((3, 800, 1333)   )]		,    segmentation_maps=[np.zeros((384, 384)   ).astype(np.floataa   ), np.zeros((384, 384)   ).astype(np.floataa   )]		,    return_tensors="pt"		,    )

             UpperCAmelCase_      			=							inputs["pixel_values"].to(_UpperCAmelCase   )
             UpperCAmelCase_      			=							[el.to(_UpperCAmelCase   ) for el in inputs["mask_labels"]]
             UpperCAmelCase_      			=							[el.to(_UpperCAmelCase   ) for el in inputs["class_labels"]]

             with torch.no_grad():
                   UpperCAmelCase_      			=							model(**_UpperCAmelCase   )

             self.assertTrue(outputs.loss is not None   )




