

'''simple docstring'''






import collections
import inspect
import unittest
from typing import Dict, List, Tuple

from transformers import MaskFormerSwinConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, torch_device
from transformers.utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import MaskFormerSwinBackbone
 from transformers.models.maskformer import MaskFormerSwinModel






class     lowercase_			:




  """simple docstring"""




  def __init__(   self     :  List[Any]      ,lowercase__     :  Optional[Any]      ,lowercase__     :  List[str]=1_3      ,lowercase__     :  Optional[int]=3_2      ,lowercase__     :  Optional[Any]=2      ,lowercase__     :  Union[str, Any]=3      ,lowercase__     :  Any=1_6      ,lowercase__     :  List[str]=[1, 2, 1]      ,lowercase__     :  List[str]=[2, 2, 4]      ,lowercase__     :  Optional[Any]=2      ,lowercase__     :  Union[str, Any]=2.0      ,lowercase__     :  Optional[int]=True      ,lowercase__     :  Tuple=0.0      ,lowercase__     :  int=0.0      ,lowercase__     :  Union[str, Any]=0.1      ,lowercase__     :  List[Any]="gelu"      ,lowercase__     :  Optional[Any]=False      ,lowercase__     :  List[Any]=True      ,lowercase__     :  Dict=0.0_2      ,lowercase__     :  Union[str, Any]=1e-5      ,lowercase__     :  List[str]=True      ,lowercase__     :  Optional[int]=None      ,lowercase__     :  str=True      ,lowercase__     :  Optional[int]=1_0      ,lowercase__     :  Dict=8      ,lowercase__     :  Optional[int]=["stage1", "stage2", "stage3"]      ,lowercase__     :  Dict=[1, 2, 3]      ,):
   __lowercase													=				parent
   __lowercase													=				batch_size
   __lowercase													=				image_size
   __lowercase													=				patch_size
   __lowercase													=				num_channels
   __lowercase													=				embed_dim
   __lowercase													=				depths
   __lowercase													=				num_heads
   __lowercase													=				window_size
   __lowercase													=				mlp_ratio
   __lowercase													=				qkv_bias
   __lowercase													=				hidden_dropout_prob
   __lowercase													=				attention_probs_dropout_prob
   __lowercase													=				drop_path_rate
   __lowercase													=				hidden_act
   __lowercase													=				use_absolute_embeddings
   __lowercase													=				patch_norm
   __lowercase													=				layer_norm_eps
   __lowercase													=				initializer_range
   __lowercase													=				is_training
   __lowercase													=				scope
   __lowercase													=				use_labels
   __lowercase													=				type_sequence_label_size
   __lowercase													=				encoder_stride
   __lowercase													=				out_features
   __lowercase													=				out_indices






  def       SCREAMING_SNAKE_CASE   (   self     :  List[Any]    ):
   __lowercase													=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

   __lowercase													=				None
   if self.use_labels:
    __lowercase													=				ids_tensor([self.batch_size]      ,self.type_sequence_label_size    )

   __lowercase													=				self.get_config()

   return config, pixel_values, labels






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[int]    ):
   return MaskFormerSwinConfig(
       image_size=self.image_size      ,patch_size=self.patch_size      ,num_channels=self.num_channels      ,embed_dim=self.embed_dim      ,depths=self.depths      ,num_heads=self.num_heads      ,window_size=self.window_size      ,mlp_ratio=self.mlp_ratio      ,qkv_bias=self.qkv_bias      ,hidden_dropout_prob=self.hidden_dropout_prob      ,attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,drop_path_rate=self.drop_path_rate      ,hidden_act=self.hidden_act      ,use_absolute_embeddings=self.use_absolute_embeddings      ,path_norm=self.patch_norm      ,layer_norm_eps=self.layer_norm_eps      ,initializer_range=self.initializer_range      ,encoder_stride=self.encoder_stride      ,out_features=self.out_features      ,out_indices=self.out_indices      ,)






  def       SCREAMING_SNAKE_CASE   (   self     :  str      ,lowercase__     :  Union[str, Any]      ,lowercase__     :  Dict      ,lowercase__     :  List[Any]    ):
   __lowercase													=				MaskFormerSwinModel(config=lowercase__    )
   model.to(lowercase__    )
   model.eval()
   __lowercase													=				model(lowercase__    )

   __lowercase													=				((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths    ) - 1))
   __lowercase													=				int(config.embed_dim * 2 ** (len(config.depths    ) - 1)    )

   self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, expected_seq_len, expected_dim)    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Union[str, Any]      ,lowercase__     :  Dict      ,lowercase__     :  Union[str, Any]      ,lowercase__     :  Union[str, Any]    ):
   __lowercase													=				MaskFormerSwinBackbone(config=lowercase__    )
   model.to(lowercase__    )
   model.eval()
   __lowercase													=				model(lowercase__    )

   # verify feature maps
   self.parent.assertEqual(len(result.feature_maps    )      ,len(config.out_features    )    )
   self.parent.assertListEqual(list(result.feature_maps[0].shape    )      ,[1_3, 1_6, 1_6, 1_6]    )

   # verify channels
   self.parent.assertEqual(len(model.channels    )      ,len(config.out_features    )    )
   self.parent.assertListEqual(model.channels      ,[1_6, 3_2, 6_4]    )

   # verify ValueError
   with self.parent.assertRaises(lowercase__    ):
    __lowercase													=				['''stem''']
    __lowercase													=				MaskFormerSwinBackbone(config=lowercase__    )






  def       SCREAMING_SNAKE_CASE   (   self     :  List[str]    ):
   __lowercase													=				self.prepare_config_and_inputs()
   __lowercase						,			__lowercase						,			__lowercase													=				config_and_inputs
   __lowercase													=				{'''pixel_values''': pixel_values}
   return config, inputs_dict






@require_torch
class     lowercase_			(lowerCamelCase__       , lowerCamelCase__       , unittest.TestCase   ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE					:						Tuple           =    (
      (
          MaskFormerSwinModel,
          MaskFormerSwinBackbone,
      )
      if is_torch_available()
      else ()
  )
  SCREAMING_SNAKE_CASE					:						Dict           =    {'feature-extraction': MaskFormerSwinModel} if is_torch_available() else {}
  SCREAMING_SNAKE_CASE					:						List[Any]           =    False
  SCREAMING_SNAKE_CASE					:						Any           =    False
  SCREAMING_SNAKE_CASE					:						int           =    False
  SCREAMING_SNAKE_CASE					:						Tuple           =    False
  SCREAMING_SNAKE_CASE					:						Dict           =    False






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[Any]    ):
   __lowercase													=				MaskFormerSwinModelTester(self    )
   __lowercase													=				ConfigTester(self      ,config_class=lowercase__      ,embed_dim=3_7    )






  @require_torch_multi_gpu
  @unittest.skip(
      reason=(
          '''`MaskFormerSwinModel` outputs `hidden_states_spatial_dimensions` which doesn\'t work well with'''
          ''' `nn.DataParallel`'''
      )    )
  def       SCREAMING_SNAKE_CASE   (   self     :  Dict    ):
   pass






  def       SCREAMING_SNAKE_CASE   (   self     :  Union[str, Any]    ):
   self.create_and_test_config_common_properties()
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   return






  def       SCREAMING_SNAKE_CASE   (   self     :  Dict    ):
   __lowercase													=				self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*lowercase__    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_backbone(*lowercase__    )






  @unittest.skip('''Swin does not use inputs_embeds'''    )
  def       SCREAMING_SNAKE_CASE   (   self     :  List[Any]    ):
   pass






  @unittest.skip('''Swin does not support feedforward chunking'''    )
  def       SCREAMING_SNAKE_CASE   (   self     :  Union[str, Any]    ):
   pass






  def       SCREAMING_SNAKE_CASE   (   self     :  Tuple    ):
   __lowercase						,			__lowercase													=				self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    __lowercase													=				model_class(lowercase__    )
    self.assertIsInstance(model.get_input_embeddings()      ,(nn.Module)    )
    __lowercase													=				model.get_output_embeddings()
    self.assertTrue(x is None or isinstance(lowercase__      ,nn.Linear    )    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Union[str, Any]    ):
   __lowercase						,			__lowercase													=				self.model_tester.prepare_config_and_inputs_for_common()

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

    __lowercase													=				['''pixel_values''']
    self.assertListEqual(arg_names[:1]      ,lowercase__    )






  @unittest.skip(reason='''MaskFormerSwin is only used as backbone and doesn\'t support output_attentions'''    )
  def       SCREAMING_SNAKE_CASE   (   self     :  List[str]    ):
   pass






  @unittest.skip(reason='''MaskFormerSwin is only used as an internal backbone'''    )
  def       SCREAMING_SNAKE_CASE   (   self     :  List[Any]    ):
   pass






  def       SCREAMING_SNAKE_CASE   (   self     :  Dict      ,lowercase__     :  int      ,lowercase__     :  Union[str, Any]      ,lowercase__     :  str      ,lowercase__     :  str    ):
   __lowercase													=				model_class(lowercase__    )
   model.to(lowercase__    )
   model.eval()

   with torch.no_grad():
    __lowercase													=				model(**self._prepare_for_class(lowercase__      ,lowercase__    )    )

   __lowercase													=				outputs.hidden_states

   __lowercase													=				getattr(
       self.model_tester      ,'''expected_num_hidden_layers'''      ,len(self.model_tester.depths    ) + 1    )
   self.assertEqual(len(lowercase__    )      ,lowercase__    )

   # Swin has a different seq_length
   __lowercase													=				(
       config.patch_size
       if isinstance(config.patch_size      ,collections.abc.Iterable    )
       else (config.patch_size, config.patch_size)
   )

   __lowercase													=				(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

   self.assertListEqual(
       list(hidden_states[0].shape[-2:]    )      ,[num_patches, self.model_tester.embed_dim]      ,)






  def       SCREAMING_SNAKE_CASE   (   self     :  List[str]    ):
   __lowercase						,			__lowercase													=				self.model_tester.prepare_config_and_inputs_for_common()

   __lowercase													=				(
       self.model_tester.image_size
       if isinstance(self.model_tester.image_size      ,collections.abc.Iterable    )
       else (self.model_tester.image_size, self.model_tester.image_size)
   )

   for model_class in self.all_model_classes:
    __lowercase													=				True
    self.check_hidden_states_output(lowercase__      ,lowercase__      ,lowercase__      ,lowercase__    )

    # check that output_hidden_states also work using config
    del inputs_dict["output_hidden_states"]
    __lowercase													=				True

    self.check_hidden_states_output(lowercase__      ,lowercase__      ,lowercase__      ,lowercase__    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[Any]    ):
   __lowercase						,			__lowercase													=				self.model_tester.prepare_config_and_inputs_for_common()
   __lowercase													=				3

   __lowercase													=				(
       self.model_tester.image_size
       if isinstance(self.model_tester.image_size      ,collections.abc.Iterable    )
       else (self.model_tester.image_size, self.model_tester.image_size)
   )
   __lowercase													=				(
       config.patch_size
       if isinstance(config.patch_size      ,collections.abc.Iterable    )
       else (config.patch_size, config.patch_size)
   )

   __lowercase													=				image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
   __lowercase													=				image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

   for model_class in self.all_model_classes:
    __lowercase													=				True
    self.check_hidden_states_output(lowercase__      ,lowercase__      ,lowercase__      ,(padded_height, padded_width)    )

    # check that output_hidden_states also work using config
    del inputs_dict["output_hidden_states"]
    __lowercase													=				True
    self.check_hidden_states_output(lowercase__      ,lowercase__      ,lowercase__      ,(padded_height, padded_width)    )






  @unittest.skip(reason='''MaskFormerSwin doesn\'t have pretrained checkpoints'''    )
  def       SCREAMING_SNAKE_CASE   (   self     :  List[Any]    ):
   pass






  @unittest.skip(reason='''This will be fixed once MaskFormerSwin is replaced by native Swin'''    )
  def       SCREAMING_SNAKE_CASE   (   self     :  int    ):
   pass






  @unittest.skip(reason='''This will be fixed once MaskFormerSwin is replaced by native Swin'''    )
  def       SCREAMING_SNAKE_CASE   (   self     :  List[str]    ):
   pass






  def       SCREAMING_SNAKE_CASE   (   self     :  str    ):
   __lowercase						,			__lowercase													=				self.model_tester.prepare_config_and_inputs_for_common()

   def set_nan_tensor_to_zero(lowercase__     :  Union[str, Any]    ):
    __lowercase													=				0
    return t

   def check_equivalence(lowercase__     :  List[str]      ,lowercase__     :  Optional[int]      ,lowercase__     :  str      ,lowercase__     :  Union[str, Any]={}    ):
    with torch.no_grad():
     __lowercase													=				model(**lowercase__      ,return_dict=lowercase__      ,**lowercase__    )
     __lowercase													=				model(**lowercase__      ,return_dict=lowercase__      ,**lowercase__    ).to_tuple()

     def recursive_check(lowercase__     :  Optional[int]      ,lowercase__     :  Optional[Any]    ):
      if isinstance(lowercase__      ,(List, Tuple)    ):
       for tuple_iterable_value, dict_iterable_value in zip(lowercase__      ,lowercase__    ):
        recursive_check(lowercase__      ,lowercase__    )
      elif isinstance(lowercase__      ,lowercase__    ):
       for tuple_iterable_value, dict_iterable_value in zip(
           tuple_object.values()      ,dict_object.values()    ):
        recursive_check(lowercase__      ,lowercase__    )
      elif tuple_object is None:
       return
      else:
       self.assertTrue(
           torch.allclose(
               set_nan_tensor_to_zero(lowercase__    )      ,set_nan_tensor_to_zero(lowercase__    )      ,atol=1e-5    )      ,msg=(
               '''Tuple and dict output are not equal. Difference:'''
               F" {torch.max(torch.abs(tuple_object - dict_object    )    )}. Tuple has `nan`:"
               F" {torch.isnan(lowercase__    ).any()} and `inf`: {torch.isinf(lowercase__    )}. Dict has"
               F" `nan`: {torch.isnan(lowercase__    ).any()} and `inf`: {torch.isinf(lowercase__    )}."
           )      ,)

     recursive_check(lowercase__      ,lowercase__    )

   for model_class in self.all_model_classes:
    __lowercase													=				model_class(lowercase__    )
    model.to(lowercase__    )
    model.eval()

    __lowercase													=				self._prepare_for_class(lowercase__      ,lowercase__    )
    __lowercase													=				self._prepare_for_class(lowercase__      ,lowercase__    )
    check_equivalence(lowercase__      ,lowercase__      ,lowercase__    )

    __lowercase													=				self._prepare_for_class(lowercase__      ,lowercase__      ,return_labels=lowercase__    )
    __lowercase													=				self._prepare_for_class(lowercase__      ,lowercase__      ,return_labels=lowercase__    )
    check_equivalence(lowercase__      ,lowercase__      ,lowercase__    )

    __lowercase													=				self._prepare_for_class(lowercase__      ,lowercase__    )
    __lowercase													=				self._prepare_for_class(lowercase__      ,lowercase__    )
    check_equivalence(lowercase__      ,lowercase__      ,lowercase__      ,{'''output_hidden_states''': True}    )

    __lowercase													=				self._prepare_for_class(lowercase__      ,lowercase__      ,return_labels=lowercase__    )
    __lowercase													=				self._prepare_for_class(lowercase__      ,lowercase__      ,return_labels=lowercase__    )
    check_equivalence(lowercase__      ,lowercase__      ,lowercase__      ,{'''output_hidden_states''': True}    )






@require_torch
class     lowercase_			(unittest.TestCase       , lowerCamelCase__   ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE					:						Union[str, Any]           =    (MaskFormerSwinBackbone,) if is_torch_available() else ()
  SCREAMING_SNAKE_CASE					:						Dict           =    MaskFormerSwinConfig






  def       SCREAMING_SNAKE_CASE   (   self     :  int    ):
   __lowercase													=				MaskFormerSwinModelTester(self    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[int]    ):
   __lowercase						,			__lowercase													=				self.model_tester.prepare_config_and_inputs_for_common()
   __lowercase													=				inputs_dict['''pixel_values'''].shape[0]

   for backbone_class in self.all_model_classes:
    __lowercase													=				backbone_class(lowercase__    )
    backbone.to(lowercase__    )
    backbone.eval()

    __lowercase													=				backbone(**lowercase__    )

    # Test default outputs and verify feature maps
    self.assertIsInstance(outputs.feature_maps      ,lowercase__    )
    self.assertTrue(len(outputs.feature_maps    ) == len(backbone.channels    )    )
    for feature_map, n_channels in zip(outputs.feature_maps      ,backbone.channels    ):
     self.assertTrue(feature_map.shape[:2]      ,(batch_size, n_channels)    )
    self.assertIsNone(outputs.hidden_states    )
    self.assertIsNone(outputs.attentions    )

    # Test output_hidden_states=True
    __lowercase													=				backbone(**lowercase__      ,output_hidden_states=lowercase__    )
    self.assertIsNotNone(outputs.hidden_states    )
    self.assertTrue(len(outputs.hidden_states    )      ,len(backbone.stage_names    )    )
    # We skip the stem layer
    for hidden_states, n_channels in zip(outputs.hidden_states[1:]      ,backbone.channels    ):
     for hidden_state in hidden_states:
      # Hidden states are in the format (batch_size, (height * width), n_channels)
      __lowercase						,			__lowercase						,			__lowercase													=				hidden_state.shape
      self.assertTrue((h_batch_size, h_n_channels)      ,(batch_size, n_channels)    )

            # Test output_attentions=True
    if self.has_attentions:
     __lowercase													=				backbone(**lowercase__      ,output_attentions=lowercase__    )
     self.assertIsNotNone(outputs.attentions    )




'''simple docstring'''






import argparse
import json
import os

import torch
from transformers.file_utils import has_file

from diffusers import UNetaDConditionModel, UNetaDModel


lowerCAmelCase__      		=			False
lowerCAmelCase__      		=			True
lowerCAmelCase__      		=			False


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

 parser.add_argument(
     '''--repo_path''',
     default=None,
     type=str,
     required=True,
     help='''The config json file corresponding to the architecture.''',
 )

 parser.add_argument('''--dump_path''', default=None, type=str, required=True, help='''Path to the output model.''')

 lowerCAmelCase__      		=			parser.parse_args()

 lowerCAmelCase__      		=			{
     '''image_size''': '''sample_size''',
     '''num_res_blocks''': '''layers_per_block''',
     '''block_channels''': '''block_out_channels''',
     '''down_blocks''': '''down_block_types''',
     '''up_blocks''': '''up_block_types''',
     '''downscale_freq_shift''': '''freq_shift''',
     '''resnet_num_groups''': '''norm_num_groups''',
     '''resnet_act_fn''': '''act_fn''',
     '''resnet_eps''': '''norm_eps''',
     '''num_head_channels''': '''attention_head_dim''',
 }

 lowerCAmelCase__      		=			{
     '''time_steps''': '''time_proj''',
     '''mid''': '''mid_block''',
     '''downsample_blocks''': '''down_blocks''',
     '''upsample_blocks''': '''up_blocks''',
 }

 lowerCAmelCase__      		=			'''''' if has_file(args.repo_path, '''config.json''') else '''unet'''

 with open(os.path.join(args.repo_path, subfolder, '''config.json'''), '''r''', encoding='''utf-8''') as reader:
  lowerCAmelCase__      		=			reader.read()
  lowerCAmelCase__      		=			json.loads(text)

 if do_only_config:
  for key in config_parameters_to_change.keys():
   config.pop(key, None)

 if has_file(args.repo_path, '''config.json'''):
  lowerCAmelCase__      		=			UNetaDModel(**config)
 else:
  lowerCAmelCase__      		=			UNetaDConditionModel if '''ldm-text2im-large-256''' in args.repo_path else UNetaDModel
  lowerCAmelCase__      		=			class_name(**config)

 if do_only_config:
  model.save_config(os.path.join(args.repo_path, subfolder))

 lowerCAmelCase__      		=			dict(model.config)

 if do_only_renaming:
  for key, value in config_parameters_to_change.items():
   if key in config:
    lowerCAmelCase__      		=			config[key]
    del config[key]

  lowerCAmelCase__      		=			[k.replace('''UNetRes''', '''''') for k in config['''down_block_types''']]
  lowerCAmelCase__      		=			[k.replace('''UNetRes''', '''''') for k in config['''up_block_types''']]

 if do_only_weights:
  lowerCAmelCase__      		=			torch.load(os.path.join(args.repo_path, subfolder, '''diffusion_pytorch_model.bin'''))

  lowerCAmelCase__      		=			{}
  for param_key, param_value in state_dict.items():
   if param_key.endswith('''.op.bias''') or param_key.endswith('''.op.weight'''):
    continue
   lowerCAmelCase__      		=			False
   for key, new_key in key_parameters_to_change.items():
    if not has_changed and param_key.split('''.''')[0] == key:
     lowerCAmelCase__      		=			param_value
     lowerCAmelCase__      		=			True
   if not has_changed:
    lowerCAmelCase__      		=			param_value

  model.load_state_dict(new_state_dict)
  model.save_pretrained(os.path.join(args.repo_path, subfolder))

