
"""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       lowerCAmelCase_		:



							'''simple docstring'''





							def __init__( self    ,				snake_case_    ,				snake_case_=2    ,				snake_case_=True    ,				snake_case_=False    ,				snake_case_=10    ,				snake_case_=3    ,				snake_case_=32 * 8    ,				snake_case_=32 * 8    ,				snake_case_=4    ,				snake_case_=64    ,				)       ->  Optional[Any]:
													__lowerCAmelCase     =  parent
													__lowerCAmelCase     =  batch_size
													__lowerCAmelCase     =  is_training
													__lowerCAmelCase     =  use_auxiliary_loss
													__lowerCAmelCase     =  num_queries
													__lowerCAmelCase     =  num_channels
													__lowerCAmelCase     =  min_size
													__lowerCAmelCase     =  max_size
													__lowerCAmelCase     =  num_labels
													__lowerCAmelCase     =  hidden_dim
													__lowerCAmelCase     =  hidden_dim




							def    A__						( self				)       ->  Tuple:
													__lowerCAmelCase     =  floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]				).to(
													    snake_case_				)

													__lowerCAmelCase     =  torch.ones([self.batch_size, self.min_size, self.max_size]    ,				device=snake_case_				)

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

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




							def    A__						( self				)       ->  Optional[int]:
													__lowerCAmelCase     =  MaskaFormerConfig(
													    hidden_size=self.hidden_dim    ,				)
													__lowerCAmelCase     =  self.num_queries
													__lowerCAmelCase     =  self.num_labels

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

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




							def    A__						( self				)       ->  Any:
													__lowerCAmelCase   ,     __lowerCAmelCase   ,     __lowerCAmelCase   ,     __lowerCAmelCase   ,     __lowerCAmelCase     =  self.prepare_config_and_inputs()
													__lowerCAmelCase     =  {"""pixel_values""": pixel_values, """pixel_mask""": pixel_mask}
													return config, inputs_dict




							def    A__						( self    ,				snake_case_    ,				snake_case_				)       ->  Dict:
													__lowerCAmelCase     =  output.encoder_hidden_states
													__lowerCAmelCase     =  output.pixel_decoder_hidden_states
													__lowerCAmelCase     =  output.transformer_decoder_hidden_states

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




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_=False				)       ->  int:
													with torch.no_grad():
																			__lowerCAmelCase     =  MaskaFormerModel(config=snake_case_				)
																			model.to(snake_case_				)
																			model.eval()

																			__lowerCAmelCase     =  model(pixel_values=snake_case_    ,				pixel_mask=snake_case_				)
																			__lowerCAmelCase     =  model(snake_case_    ,				output_hidden_states=snake_case_				)

													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(snake_case_    ,				snake_case_				)




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  Optional[Any]:
													__lowerCAmelCase     =  MaskaFormerForUniversalSegmentation(config=snake_case_				)
													model.to(snake_case_				)
													model.eval()

													def comm_check_on_output(snake_case_				):
																			# 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():
																			__lowerCAmelCase     =  model(pixel_values=snake_case_    ,				pixel_mask=snake_case_				)
																			__lowerCAmelCase     =  model(snake_case_				)

																			comm_check_on_output(snake_case_				)

																			__lowerCAmelCase     =  model(
																			    pixel_values=snake_case_    ,				pixel_mask=snake_case_    ,				mask_labels=snake_case_    ,				class_labels=snake_case_				)

													comm_check_on_output(snake_case_				)

													self.parent.assertTrue(result.loss is not None				)
													self.parent.assertEqual(result.loss.shape    ,				torch.Size([1]				)				)
@require_torch
class       lowerCAmelCase_		(				A__				,				A__				,				unittest.TestCase ):



							'''simple docstring'''





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

							_snake_case						      =		False
							_snake_case						      =		False
							_snake_case						      =		False
							_snake_case						      =		False




							def    A__						( self				)       ->  Optional[Any]:
													__lowerCAmelCase     =  MaskaFormerModelTester(self				)
													__lowerCAmelCase     =  ConfigTester(self    ,				config_class=snake_case_    ,				has_text_modality=snake_case_				)




							def    A__						( self				)       ->  List[str]:
													self.config_tester.run_common_tests()




							def    A__						( self				)       ->  List[str]:
													__lowerCAmelCase   ,     __lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs_for_common()
													self.model_tester.create_and_check_maskaformer_model(snake_case_    ,				**snake_case_    ,				output_hidden_states=snake_case_				)




							def    A__						( self				)       ->  Any:
													__lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*snake_case_				)




							@unittest.skip(reason="""Mask2Former does not use inputs_embeds"""				)
							def    A__						( self				)       ->  List[Any]:
													pass




							@unittest.skip(reason="""Mask2Former does not have a get_input_embeddings method"""				)
							def    A__						( self				)       ->  str:
													pass




							@unittest.skip(reason="""Mask2Former is not a generative model"""				)
							def    A__						( self				)       ->  Optional[int]:
													pass




							@unittest.skip(reason="""Mask2Former does not use token embeddings"""				)
							def    A__						( self				)       ->  List[str]:
													pass




							@require_torch_multi_gpu
							@unittest.skip(
							    reason="""Mask2Former has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""				)
							def    A__						( self				)       ->  int:
													pass




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




							def    A__						( self				)       ->  List[Any]:
													__lowerCAmelCase   ,     __lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs_for_common()

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

																			__lowerCAmelCase     =  ["""pixel_values"""]
																			self.assertListEqual(arg_names[:1]    ,				snake_case_				)




							@slow
							def    A__						( self				)       ->  Optional[int]:
													for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
																			__lowerCAmelCase     =  MaskaFormerModel.from_pretrained(snake_case_				)
																			self.assertIsNotNone(snake_case_				)




							def    A__						( self				)       ->  Optional[Any]:
													__lowerCAmelCase     =  (self.model_tester.min_size,) * 2
													__lowerCAmelCase     =  {
													    """pixel_values""": torch.randn((2, 3, *size)    ,				device=snake_case_				),
													    """mask_labels""": torch.randn((2, 10, *size)    ,				device=snake_case_				),
													    """class_labels""": torch.zeros(2    ,				10    ,				device=snake_case_				).long(),
													}
													__lowerCAmelCase     =  self.model_tester.get_config()

													__lowerCAmelCase     =  MaskaFormerForUniversalSegmentation(snake_case_				).to(snake_case_				)
													__lowerCAmelCase     =  model(**snake_case_				)
													self.assertTrue(outputs.loss is not None				)




							def    A__						( self				)       ->  List[str]:
													__lowerCAmelCase   ,     __lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs_for_common()
													self.model_tester.create_and_check_maskaformer_model(snake_case_    ,				**snake_case_    ,				output_hidden_states=snake_case_				)




							def    A__						( self				)       ->  str:
													__lowerCAmelCase   ,     __lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs_for_common()

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




							def    A__						( self				)       ->  int:
													if not self.model_tester.is_training:
																			return

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

													__lowerCAmelCase     =  model_class(snake_case_				)
													model.to(snake_case_				)
													model.train()

													__lowerCAmelCase     =  model(snake_case_    ,				mask_labels=snake_case_    ,				class_labels=snake_case_				).loss
													loss.backward()




							def    A__						( self				)       ->  Optional[int]:
													__lowerCAmelCase     =  self.all_model_classes[1]
													__lowerCAmelCase   ,     __lowerCAmelCase   ,     __lowerCAmelCase   ,     __lowerCAmelCase   ,     __lowerCAmelCase     =  self.model_tester.prepare_config_and_inputs()
													__lowerCAmelCase     =  True
													__lowerCAmelCase     =  True

													__lowerCAmelCase     =  model_class(snake_case_				).to(snake_case_				)
													model.train()

													__lowerCAmelCase     =  model(snake_case_    ,				mask_labels=snake_case_    ,				class_labels=snake_case_				)

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

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

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

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

													outputs.loss.backward(retain_graph=snake_case_				)

													self.assertIsNotNone(encoder_hidden_states.grad				)
													self.assertIsNotNone(pixel_decoder_hidden_states.grad				)
													self.assertIsNotNone(transformer_decoder_hidden_states.grad				)
													self.assertIsNotNone(attentions.grad				)
SCREAMING_SNAKE_CASE_					     =	1E-4

def  lowercase     ():
						__lowerCAmelCase     =  Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""						)
						return image



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



							'''simple docstring'''





							@cached_property
							def    A__						( self				)       ->  Union[str, Any]:
													return "facebook/mask2former-swin-small-coco-instance"




							@cached_property
							def    A__						( self				)       ->  List[Any]:
													return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints				) if is_vision_available() else None




							def    A__						( self				)       ->  List[str]:
													__lowerCAmelCase     =  MaskaFormerModel.from_pretrained(self.model_checkpoints				).to(snake_case_				)
													__lowerCAmelCase     =  self.default_image_processor
													__lowerCAmelCase     =  prepare_img()
													__lowerCAmelCase     =  image_processor(snake_case_    ,				return_tensors="""pt"""				).to(snake_case_				)
													__lowerCAmelCase     =  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(snake_case_    ,				(1, 3, 384, 384)				)

													with torch.no_grad():
																			__lowerCAmelCase     =  model(**snake_case_				)

													__lowerCAmelCase     =  torch.tensor(
													    [[-0.2_790, -1.0_717, -1.1_668], [-0.5_128, -0.3_128, -0.4_987], [-0.5_832, 0.1_971, -0.0_197]]				).to(snake_case_				)
													self.assertTrue(
													    torch.allclose(
													        outputs.encoder_last_hidden_state[0, 0, :3, :3]    ,				snake_case_    ,				atol=snake_case_				)				)

													__lowerCAmelCase     =  torch.tensor(
													    [[0.8_973, 1.1_847, 1.1_776], [1.1_934, 1.5_040, 1.5_128], [1.1_153, 1.4_486, 1.4_951]]				).to(snake_case_				)
													self.assertTrue(
													    torch.allclose(
													        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]    ,				snake_case_    ,				atol=snake_case_				)				)

													__lowerCAmelCase     =  torch.tensor(
													    [[2.1_152, 1.7_000, -0.8_603], [1.5_808, 1.8_004, -0.9_353], [1.6_043, 1.7_495, -0.5_999]]				).to(snake_case_				)
													self.assertTrue(
													    torch.allclose(
													        outputs.transformer_decoder_last_hidden_state[0, :3, :3]    ,				snake_case_    ,				atol=snake_case_				)				)




							def    A__						( self				)       ->  Any:
													__lowerCAmelCase     =  MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints				).to(snake_case_				).eval()
													__lowerCAmelCase     =  self.default_image_processor
													__lowerCAmelCase     =  prepare_img()
													__lowerCAmelCase     =  image_processor(snake_case_    ,				return_tensors="""pt"""				).to(snake_case_				)
													__lowerCAmelCase     =  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(snake_case_    ,				(1, 3, 384, 384)				)

													with torch.no_grad():
																			__lowerCAmelCase     =  model(**snake_case_				)
													# masks_queries_logits
													__lowerCAmelCase     =  outputs.masks_queries_logits
													self.assertEqual(
													    masks_queries_logits.shape    ,				(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)				)
													__lowerCAmelCase     =  [
													    [-8.7_839, -9.0_056, -8.8_121],
													    [-7.4_104, -7.0_313, -6.5_401],
													    [-6.6_105, -6.3_427, -6.4_675],
													]
													__lowerCAmelCase     =  torch.tensor(snake_case_				).to(snake_case_				)
													self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]    ,				snake_case_    ,				atol=snake_case_				)				)
													# class_queries_logits
													__lowerCAmelCase     =  outputs.class_queries_logits
													self.assertEqual(class_queries_logits.shape    ,				(1, model.config.num_queries, model.config.num_labels + 1)				)
													__lowerCAmelCase     =  torch.tensor(
													    [
													        [1.8_324, -8.0_835, -4.1_922],
													        [0.8_450, -9.0_050, -3.6_053],
													        [0.3_045, -7.7_293, -3.0_275],
													    ]				).to(snake_case_				)
													self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]    ,				snake_case_    ,				atol=snake_case_				)				)




							def    A__						( self				)       ->  Tuple:
													__lowerCAmelCase     =  MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints				).to(snake_case_				).eval()
													__lowerCAmelCase     =  self.default_image_processor

													__lowerCAmelCase     =  image_processor(
													    [np.zeros((3, 800, 1_333)				), np.zeros((3, 800, 1_333)				)]    ,				segmentation_maps=[np.zeros((384, 384)				).astype(np.floataa				), np.zeros((384, 384)				).astype(np.floataa				)]    ,				return_tensors="""pt"""    ,				)

													__lowerCAmelCase     =  inputs["""pixel_values"""].to(snake_case_				)
													__lowerCAmelCase     =  [el.to(snake_case_				) for el in inputs["""mask_labels"""]]
													__lowerCAmelCase     =  [el.to(snake_case_				) for el in inputs["""class_labels"""]]

													with torch.no_grad():
																			__lowerCAmelCase     =  model(**snake_case_				)

													self.assertTrue(outputs.loss is not None				)


"""simple docstring"""
import operator as op


SCREAMING_SNAKE_CASE_					     =	'''scaler.pt'''
SCREAMING_SNAKE_CASE_					     =	'''pytorch_model'''
SCREAMING_SNAKE_CASE_					     =	'''random_states'''
SCREAMING_SNAKE_CASE_					     =	'''optimizer'''
SCREAMING_SNAKE_CASE_					     =	'''scheduler'''
SCREAMING_SNAKE_CASE_					     =	'''pytorch_model.bin'''
SCREAMING_SNAKE_CASE_					     =	'''pytorch_model.bin.index.json'''
SCREAMING_SNAKE_CASE_					     =	'''model.safetensors'''
SCREAMING_SNAKE_CASE_					     =	'''model.safetensors.index.json'''
SCREAMING_SNAKE_CASE_					     =	'''1.10.2'''
SCREAMING_SNAKE_CASE_					     =	'''py38'''
SCREAMING_SNAKE_CASE_					     =	'''4.17.0'''
SCREAMING_SNAKE_CASE_					     =	['''ml.p3.16xlarge''', '''ml.p3dn.24xlarge''', '''ml.p4dn.24xlarge''']
SCREAMING_SNAKE_CASE_					     =	['''FULL_SHARD''', '''SHARD_GRAD_OP''', '''NO_SHARD''', '''HYBRID_SHARD''', '''HYBRID_SHARD_ZERO2''']
SCREAMING_SNAKE_CASE_					     =	['''TRANSFORMER_BASED_WRAP''', '''SIZE_BASED_WRAP''', '''NO_WRAP''']
SCREAMING_SNAKE_CASE_					     =	['''BACKWARD_PRE''', '''BACKWARD_POST''', '''NO_PREFETCH''']
SCREAMING_SNAKE_CASE_					     =	['''FULL_STATE_DICT''', '''LOCAL_STATE_DICT''', '''SHARDED_STATE_DICT''']
SCREAMING_SNAKE_CASE_					     =	'''2.0.1'''
SCREAMING_SNAKE_CASE_					     =	['''pdsh''', '''standard''', '''openmpi''', '''mvapich''']
SCREAMING_SNAKE_CASE_					     =	['''default''', '''reduce-overhead''', '''max-autotune''']

SCREAMING_SNAKE_CASE_					     =	{'''>''': op.gt, '''>=''': op.ge, '''==''': op.eq, '''!=''': op.ne, '''<=''': op.le, '''<''': op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
SCREAMING_SNAKE_CASE_					     =	[
    '''nnodes''',
    '''nproc_per_node''',
    '''rdzv_backend''',
    '''rdzv_endpoint''',
    '''rdzv_id''',
    '''rdzv_conf''',
    '''standalone''',
    '''max_restarts''',
    '''monitor_interval''',
    '''start_method''',
    '''role''',
    '''module''',
    '''m''',
    '''no_python''',
    '''run_path''',
    '''log_dir''',
    '''r''',
    '''redirects''',
    '''t''',
    '''tee''',
    '''node_rank''',
    '''master_addr''',
    '''master_port''',
]

SCREAMING_SNAKE_CASE_					     =	['''DEEPSPEED''', '''MULTI_GPU''', '''FSDP''', '''MEGATRON_LM''']
SCREAMING_SNAKE_CASE_					     =	['''DEEPSPEED''', '''MULTI_XPU''', '''FSDP''']
