
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 	UpperCAmelCase_ :
	"""simple docstring"""

	def __init__(     self:     List[Any]							,   _UpperCAmelCase:     Union[str, Any]							,   _UpperCAmelCase:     Optional[Any]=2							,   _UpperCAmelCase:     Dict=True							,   _UpperCAmelCase:     Any=False							,   _UpperCAmelCase:     Dict=10							,   _UpperCAmelCase:     Optional[Any]=3							,   _UpperCAmelCase:     Dict=32 * 8							,   _UpperCAmelCase:     Dict=32 * 8							,   _UpperCAmelCase:     Optional[int]=4							,   _UpperCAmelCase:     Any=64							,   ):
								_lowerCAmelCase :str      	=						parent
								_lowerCAmelCase :Tuple      	=						batch_size
								_lowerCAmelCase :str      	=						is_training
								_lowerCAmelCase :Tuple      	=						use_auxiliary_loss
								_lowerCAmelCase :Union[str, Any]      	=						num_queries
								_lowerCAmelCase :int      	=						num_channels
								_lowerCAmelCase :Union[str, Any]      	=						min_size
								_lowerCAmelCase :List[Any]      	=						max_size
								_lowerCAmelCase :int      	=						num_labels
								_lowerCAmelCase :Optional[int]      	=						hidden_dim
								_lowerCAmelCase :Union[str, Any]      	=						hidden_dim







	def       SCREAMING_SNAKE_CASE__     (     self:     List[Any]			):
								_lowerCAmelCase :Union[str, Any]      	=						floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]			).to(
								    __a			)

								_lowerCAmelCase :Union[str, Any]      	=						torch.ones([self.batch_size, self.min_size, self.max_size]							,   device=__a			)

								_lowerCAmelCase :Tuple      	=						(
								    torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]							,   device=__a			) > 0.5
								).float()
								_lowerCAmelCase :Any      	=						(torch.rand((self.batch_size, self.num_labels)							,   device=__a			) > 0.5).long()

								_lowerCAmelCase :Any      	=						self.get_config()
								return config, pixel_values, pixel_mask, mask_labels, class_labels







	def       SCREAMING_SNAKE_CASE__     (     self:     Optional[int]			):
								_lowerCAmelCase :Union[str, Any]      	=						MaskaFormerConfig(
								    hidden_size=self.hidden_dim							,   )
								_lowerCAmelCase :Dict      	=						self.num_queries
								_lowerCAmelCase :Any      	=						self.num_labels

								_lowerCAmelCase :Tuple      	=						[1, 1, 1, 1]
								_lowerCAmelCase :int      	=						self.num_channels

								_lowerCAmelCase :Tuple      	=						64
								_lowerCAmelCase :Dict      	=						128
								_lowerCAmelCase :Optional[int]      	=						self.hidden_dim
								_lowerCAmelCase :int      	=						self.hidden_dim
								_lowerCAmelCase :Union[str, Any]      	=						self.hidden_dim
								return config







	def       SCREAMING_SNAKE_CASE__     (     self:     Optional[int]			):
								_lowerCAmelCase :int      	=						self.prepare_config_and_inputs()
								_lowerCAmelCase :List[str]      	=						{"""pixel_values""": pixel_values, """pixel_mask""": pixel_mask}
								return config, inputs_dict







	def       SCREAMING_SNAKE_CASE__     (     self:     Any							,   _UpperCAmelCase:     Union[str, Any]							,   _UpperCAmelCase:     str			):
								_lowerCAmelCase :Any      	=						output.encoder_hidden_states
								_lowerCAmelCase :List[str]      	=						output.pixel_decoder_hidden_states
								_lowerCAmelCase :Tuple      	=						output.transformer_decoder_hidden_states

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







	def       SCREAMING_SNAKE_CASE__     (     self:     Tuple							,   _UpperCAmelCase:     Optional[Any]							,   _UpperCAmelCase:     Dict							,   _UpperCAmelCase:     Dict							,   _UpperCAmelCase:     List[Any]=False			):
								with torch.no_grad():
															_lowerCAmelCase :Any      	=						MaskaFormerModel(config=__a			)
															model.to(__a			)
															model.eval()

															_lowerCAmelCase :Tuple      	=						model(pixel_values=__a							,   pixel_mask=__a			)
															_lowerCAmelCase :List[Any]      	=						model(__a							,   output_hidden_states=__a			)

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







	def       SCREAMING_SNAKE_CASE__     (     self:     str							,   _UpperCAmelCase:     Tuple							,   _UpperCAmelCase:     Optional[Any]							,   _UpperCAmelCase:     List[str]							,   _UpperCAmelCase:     Dict							,   _UpperCAmelCase:     Dict			):
								_lowerCAmelCase :Optional[Any]      	=						MaskaFormerForUniversalSegmentation(config=__a			)
								model.to(__a			)
								model.eval()

								def comm_check_on_output(_UpperCAmelCase:     str			):
															# 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 :Any      	=						model(pixel_values=__a							,   pixel_mask=__a			)
															_lowerCAmelCase :Optional[int]      	=						model(__a			)

															comm_check_on_output(__a			)

															_lowerCAmelCase :List[str]      	=						model(
															    pixel_values=__a							,   pixel_mask=__a							,   mask_labels=__a							,   class_labels=__a			)

								comm_check_on_output(__a			)

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





@require_torch
class 	UpperCAmelCase_ (a__			,  a__			,  unittest.TestCase							):
	"""simple docstring"""

	lowerCamelCase     :       Union[str, Any]            =			(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
	lowerCamelCase     :       Optional[int]            =			{"""feature-extraction""": MaskaFormerModel} if is_torch_available() else {}

	lowerCamelCase     :       Optional[Any]            =			False
	lowerCamelCase     :       Optional[int]            =			False
	lowerCamelCase     :       Dict            =			False
	lowerCamelCase     :       Optional[int]            =			False







	def       SCREAMING_SNAKE_CASE__     (     self:     Any			):
								_lowerCAmelCase :Optional[Any]      	=						MaskaFormerModelTester(self			)
								_lowerCAmelCase :List[Any]      	=						ConfigTester(self							,   config_class=__a							,   has_text_modality=__a			)







	def       SCREAMING_SNAKE_CASE__     (     self:     Any			):
								self.config_tester.run_common_tests()







	def       SCREAMING_SNAKE_CASE__     (     self:     Tuple			):
								_lowerCAmelCase :int      	=						self.model_tester.prepare_config_and_inputs_for_common()
								self.model_tester.create_and_check_maskaformer_model(__a							,   **__a							,   output_hidden_states=__a			)







	def       SCREAMING_SNAKE_CASE__     (     self:     Any			):
								_lowerCAmelCase :Any      	=						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*__a			)







	@unittest.skip(reason='Mask2Former does not use inputs_embeds'			)
	def       SCREAMING_SNAKE_CASE__     (     self:     int			):
								pass







	@unittest.skip(reason='Mask2Former does not have a get_input_embeddings method'			)
	def       SCREAMING_SNAKE_CASE__     (     self:     int			):
								pass







	@unittest.skip(reason='Mask2Former is not a generative model'			)
	def       SCREAMING_SNAKE_CASE__     (     self:     Dict			):
								pass







	@unittest.skip(reason='Mask2Former does not use token embeddings'			)
	def       SCREAMING_SNAKE_CASE__     (     self:     Tuple			):
								pass







	@require_torch_multi_gpu
	@unittest.skip(
	    reason='Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'			)
	def       SCREAMING_SNAKE_CASE__     (     self:     str			):
								pass







	@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'			)
	def       SCREAMING_SNAKE_CASE__     (     self:     int			):
								pass







	def       SCREAMING_SNAKE_CASE__     (     self:     Optional[int]			):
								_lowerCAmelCase :Optional[int]      	=						self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
															_lowerCAmelCase :int      	=						model_class(__a			)
															_lowerCAmelCase :List[Any]      	=						inspect.signature(model.forward			)
															# signature.parameters is an OrderedDict => so arg_names order is deterministic
															_lowerCAmelCase :List[str]      	=						[*signature.parameters.keys()]

															_lowerCAmelCase :Tuple      	=						["""pixel_values"""]
															self.assertListEqual(arg_names[:1]							,   __a			)







	@slow
	def       SCREAMING_SNAKE_CASE__     (     self:     List[Any]			):
								for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
															_lowerCAmelCase :List[Any]      	=						MaskaFormerModel.from_pretrained(__a			)
															self.assertIsNotNone(__a			)







	def       SCREAMING_SNAKE_CASE__     (     self:     int			):
								_lowerCAmelCase :Any      	=						(self.model_tester.min_size,) * 2
								_lowerCAmelCase :List[Any]      	=						{
								    """pixel_values""": torch.randn((2, 3, *size)							,   device=__a			),
								    """mask_labels""": torch.randn((2, 10, *size)							,   device=__a			),
								    """class_labels""": torch.zeros(2							,   10							,   device=__a			).long(),
								}
								_lowerCAmelCase :List[Any]      	=						self.model_tester.get_config()

								_lowerCAmelCase :Optional[int]      	=						MaskaFormerForUniversalSegmentation(__a			).to(__a			)
								_lowerCAmelCase :Tuple      	=						model(**__a			)
								self.assertTrue(outputs.loss is not None			)







	def       SCREAMING_SNAKE_CASE__     (     self:     str			):
								_lowerCAmelCase :int      	=						self.model_tester.prepare_config_and_inputs_for_common()
								self.model_tester.create_and_check_maskaformer_model(__a							,   **__a							,   output_hidden_states=__a			)







	def       SCREAMING_SNAKE_CASE__     (     self:     Union[str, Any]			):
								_lowerCAmelCase :int      	=						self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
															_lowerCAmelCase :Tuple      	=						model_class(__a			).to(__a			)
															_lowerCAmelCase :List[str]      	=						model(**__a							,   output_attentions=__a			)
															self.assertTrue(outputs.attentions is not None			)







	def       SCREAMING_SNAKE_CASE__     (     self:     Optional[int]			):
								if not self.model_tester.is_training:
															return

								_lowerCAmelCase :Optional[Any]      	=						self.all_model_classes[1]
								_lowerCAmelCase :Optional[int]      	=						self.model_tester.prepare_config_and_inputs()

								_lowerCAmelCase :Optional[int]      	=						model_class(__a			)
								model.to(__a			)
								model.train()

								_lowerCAmelCase :int      	=						model(__a							,   mask_labels=__a							,   class_labels=__a			).loss
								loss.backward()







	def       SCREAMING_SNAKE_CASE__     (     self:     Any			):
								_lowerCAmelCase :int      	=						self.all_model_classes[1]
								_lowerCAmelCase :Tuple      	=						self.model_tester.prepare_config_and_inputs()
								_lowerCAmelCase :Tuple      	=						True
								_lowerCAmelCase :int      	=						True

								_lowerCAmelCase :List[str]      	=						model_class(__a			).to(__a			)
								model.train()

								_lowerCAmelCase :List[Any]      	=						model(__a							,   mask_labels=__a							,   class_labels=__a			)

								_lowerCAmelCase :Optional[int]      	=						outputs.encoder_hidden_states[0]
								encoder_hidden_states.retain_grad()

								_lowerCAmelCase :Optional[int]      	=						outputs.pixel_decoder_hidden_states[0]
								pixel_decoder_hidden_states.retain_grad()

								_lowerCAmelCase :Optional[int]      	=						outputs.transformer_decoder_hidden_states[0]
								transformer_decoder_hidden_states.retain_grad()

								_lowerCAmelCase :Optional[int]      	=						outputs.attentions[0]
								attentions.retain_grad()

								outputs.loss.backward(retain_graph=__a			)

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





a											=							1E-4







def   UpperCamelCase_(					):


							"""simple docstring"""





							_lowerCAmelCase :int      	=						Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
							return image





@require_vision
@slow
class 	UpperCAmelCase_ (unittest.TestCase							):
	"""simple docstring"""

	@cached_property
	def       SCREAMING_SNAKE_CASE__     (     self:     Any			):
								return "facebook/mask2former-swin-small-coco-instance"







	@cached_property
	def       SCREAMING_SNAKE_CASE__     (     self:     List[str]			):
								return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints			) if is_vision_available() else None







	def       SCREAMING_SNAKE_CASE__     (     self:     List[Any]			):
								_lowerCAmelCase :int      	=						MaskaFormerModel.from_pretrained(self.model_checkpoints			).to(__a			)
								_lowerCAmelCase :List[Any]      	=						self.default_image_processor
								_lowerCAmelCase :Optional[Any]      	=						prepare_img()
								_lowerCAmelCase :Optional[Any]      	=						image_processor(__a							,   return_tensors='pt'			).to(__a			)
								_lowerCAmelCase :int      	=						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(__a							,   (1, 3, 384, 384)			)

								with torch.no_grad():
															_lowerCAmelCase :Union[str, Any]      	=						model(**__a			)

								_lowerCAmelCase :List[Any]      	=						torch.tensor(
								    [[-0.2_7_9_0, -1.0_7_1_7, -1.1_6_6_8], [-0.5_1_2_8, -0.3_1_2_8, -0.4_9_8_7], [-0.5_8_3_2, 0.1_9_7_1, -0.0_1_9_7]]			).to(__a			)
								self.assertTrue(
								    torch.allclose(
								        outputs.encoder_last_hidden_state[0, 0, :3, :3]							,   __a							,   atol=__a			)			)

								_lowerCAmelCase :List[str]      	=						torch.tensor(
								    [[0.8_9_7_3, 1.1_8_4_7, 1.1_7_7_6], [1.1_9_3_4, 1.5_0_4_0, 1.5_1_2_8], [1.1_1_5_3, 1.4_4_8_6, 1.4_9_5_1]]			).to(__a			)
								self.assertTrue(
								    torch.allclose(
								        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]							,   __a							,   atol=__a			)			)

								_lowerCAmelCase :Optional[Any]      	=						torch.tensor(
								    [[2.1_1_5_2, 1.7_0_0_0, -0.8_6_0_3], [1.5_8_0_8, 1.8_0_0_4, -0.9_3_5_3], [1.6_0_4_3, 1.7_4_9_5, -0.5_9_9_9]]			).to(__a			)
								self.assertTrue(
								    torch.allclose(
								        outputs.transformer_decoder_last_hidden_state[0, :3, :3]							,   __a							,   atol=__a			)			)







	def       SCREAMING_SNAKE_CASE__     (     self:     List[str]			):
								_lowerCAmelCase :List[Any]      	=						MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints			).to(__a			).eval()
								_lowerCAmelCase :Optional[Any]      	=						self.default_image_processor
								_lowerCAmelCase :List[Any]      	=						prepare_img()
								_lowerCAmelCase :int      	=						image_processor(__a							,   return_tensors='pt'			).to(__a			)
								_lowerCAmelCase :Optional[Any]      	=						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(__a							,   (1, 3, 384, 384)			)

								with torch.no_grad():
															_lowerCAmelCase :List[Any]      	=						model(**__a			)
								# masks_queries_logits
								_lowerCAmelCase :str      	=						outputs.masks_queries_logits
								self.assertEqual(
								    masks_queries_logits.shape							,   (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)			)
								_lowerCAmelCase :str      	=						[
								    [-8.7_8_3_9, -9.0_0_5_6, -8.8_1_2_1],
								    [-7.4_1_0_4, -7.0_3_1_3, -6.5_4_0_1],
								    [-6.6_1_0_5, -6.3_4_2_7, -6.4_6_7_5],
								]
								_lowerCAmelCase :str      	=						torch.tensor(__a			).to(__a			)
								self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]							,   __a							,   atol=__a			)			)
								# class_queries_logits
								_lowerCAmelCase :str      	=						outputs.class_queries_logits
								self.assertEqual(class_queries_logits.shape							,   (1, model.config.num_queries, model.config.num_labels + 1)			)
								_lowerCAmelCase :Optional[Any]      	=						torch.tensor(
								    [
								        [1.8_3_2_4, -8.0_8_3_5, -4.1_9_2_2],
								        [0.8_4_5_0, -9.0_0_5_0, -3.6_0_5_3],
								        [0.3_0_4_5, -7.7_2_9_3, -3.0_2_7_5],
								    ]			).to(__a			)
								self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]							,   __a							,   atol=__a			)			)







	def       SCREAMING_SNAKE_CASE__     (     self:     Any			):
								_lowerCAmelCase :Union[str, Any]      	=						MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints			).to(__a			).eval()
								_lowerCAmelCase :List[Any]      	=						self.default_image_processor

								_lowerCAmelCase :Dict      	=						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'							,   )

								_lowerCAmelCase :List[Any]      	=						inputs["""pixel_values"""].to(__a			)
								_lowerCAmelCase :Tuple      	=						[el.to(__a			) for el in inputs["""mask_labels"""]]
								_lowerCAmelCase :Dict      	=						[el.to(__a			) for el in inputs["""class_labels"""]]

								with torch.no_grad():
															_lowerCAmelCase :Dict      	=						model(**__a			)

								self.assertTrue(outputs.loss is not None			)





'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowercase      :   Optional[int]      =       {
    """configuration_pegasus_x""": ["""PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP""", """PegasusXConfig"""],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase      :   int      =       [
          """PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """PegasusXForConditionalGeneration""",
          """PegasusXModel""",
          """PegasusXPreTrainedModel""",
      ]


if TYPE_CHECKING:
      from .configuration_pegasus_x import PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusXConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_pegasus_x import (
                PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST,
                PegasusXForConditionalGeneration,
                PegasusXModel,
                PegasusXPreTrainedModel,
            )


else:
      import sys

      lowercase      :   Optional[Any]      =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



