


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   ,    ):

						'''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 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]  ).to(
						    snake_case  )

						UpperCamelCase__       =							torch.ones([self.batch_size, self.min_size, self.max_size]   ,    device=snake_case  )

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

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




			def 					snake_case__					(				self  ):

						'''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 					snake_case__					(				self  ):

						'''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 					snake_case__					(				self   ,    snake_case   ,    snake_case  ):

						'''simple docstring'''


						UpperCamelCase__       =							output.encoder_hidden_states
						UpperCamelCase__       =							output.pixel_decoder_hidden_states
						UpperCamelCase__       =							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 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case=False  ):

						'''simple docstring'''


						with torch.no_grad():
									UpperCamelCase__       =							MaskaFormerModel(config=snake_case  )
									model.to(snake_case  )
									model.eval()

									UpperCamelCase__       =							model(pixel_values=snake_case   ,    pixel_mask=snake_case  )
									UpperCamelCase__       =							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 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case  ):

						'''simple docstring'''


						UpperCamelCase__       =							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():
									UpperCamelCase__       =							model(pixel_values=snake_case   ,    pixel_mask=snake_case  )
									UpperCamelCase__       =							model(snake_case  )

									comm_check_on_output(snake_case  )

									UpperCamelCase__       =							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__						(				UpperCAmelCase       ,						UpperCAmelCase       ,						unittest.TestCase   ):






			"""simple docstring"""






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

			_UpperCamelCase		:       Union[str, Any] 			=		False
			_UpperCamelCase		:       Union[str, Any] 			=		False
			_UpperCamelCase		:       int 			=		False
			_UpperCamelCase		:       Optional[int] 			=		False




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							MaskaFormerModelTester(self  )
						UpperCamelCase__       =							ConfigTester(self   ,    config_class=snake_case   ,    has_text_modality=snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						self.config_tester.run_common_tests()




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__,   UpperCamelCase__       =							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 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							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 					snake_case__					(				self  ):

						'''simple docstring'''


						pass




			@unittest.skip(reason="Mask2Former does not have a get_input_embeddings method"  )
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						pass




			@unittest.skip(reason="Mask2Former is not a generative model"  )
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						pass




			@unittest.skip(reason="Mask2Former does not use token embeddings"  )
			def 					snake_case__					(				self  ):

						'''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 					snake_case__					(				self  ):

						'''simple docstring'''


						pass




			@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."  )
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						pass




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__,   UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									UpperCamelCase__       =							model_class(snake_case  )
									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]   ,    snake_case  )




			@slow
			def 					snake_case__					(				self  ):

						'''simple docstring'''


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




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							(self.model_tester.min_size,) * 2
						UpperCamelCase__       =							{
						    "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(),
						}
						UpperCamelCase__       =							self.model_tester.get_config()

						UpperCamelCase__       =							MaskaFormerForUniversalSegmentation(snake_case  ).to(snake_case  )
						UpperCamelCase__       =							model(**snake_case  )
						self.assertTrue(outputs.loss is not None  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__,   UpperCamelCase__       =							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 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__,   UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs_for_common()

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




			def 					snake_case__					(				self  ):

						'''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(snake_case  )
						model.to(snake_case  )
						model.train()

						UpperCamelCase__       =							model(snake_case   ,    mask_labels=snake_case   ,    class_labels=snake_case  ).loss
						loss.backward()







			def 					snake_case__					(				self  ):

						'''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(snake_case  ).to(snake_case  )
						model.train()

						UpperCamelCase__       =							model(snake_case   ,    mask_labels=snake_case   ,    class_labels=snake_case  )

						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=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  )



__UpperCamelCase     		=							1E-4
def 							UpperCamelCase_(					)->		str:
			UpperCamelCase__       =							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
			return image







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






			"""simple docstring"""






			@cached_property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


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




			@cached_property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


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




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							MaskaFormerModel.from_pretrained(self.model_checkpoints  ).to(snake_case  )
						UpperCamelCase__       =							self.default_image_processor
						UpperCamelCase__       =							prepare_img()
						UpperCamelCase__       =							image_processor(snake_case   ,    return_tensors="pt"  ).to(snake_case  )
						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(snake_case   ,    (1, 3, 384, 384)  )

						with torch.no_grad():
									UpperCamelCase__       =							model(**snake_case  )

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

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

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




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints  ).to(snake_case  ).eval()
						UpperCamelCase__       =							self.default_image_processor
						UpperCamelCase__       =							prepare_img()
						UpperCamelCase__       =							image_processor(snake_case   ,    return_tensors="pt"  ).to(snake_case  )
						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(snake_case   ,    (1, 3, 384, 384)  )

						with torch.no_grad():
									UpperCamelCase__       =							model(**snake_case  )
						# 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(snake_case  ).to(snake_case  )
						self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]   ,    snake_case   ,    atol=snake_case  )  )
						# 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(snake_case  )
						self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]   ,    snake_case   ,    atol=snake_case  )  )







			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints  ).to(snake_case  ).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(snake_case  )
						UpperCamelCase__       =							[el.to(snake_case  ) for el in inputs["mask_labels"]]
						UpperCamelCase__       =							[el.to(snake_case  ) for el in inputs["class_labels"]]

						with torch.no_grad():
									UpperCamelCase__       =							model(**snake_case  )

						self.assertTrue(outputs.loss is not None  )






from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__UpperCamelCase     		=							logging.get_logger(__name__)

__UpperCamelCase     		=							{
    'facebook/xmod-base': 'https://huggingface.co/facebook/xmod-base/resolve/main/config.json',
    'facebook/xmod-large-prenorm': 'https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json',
    'facebook/xmod-base-13-125k': 'https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json',
    'facebook/xmod-base-30-125k': 'https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json',
    'facebook/xmod-base-30-195k': 'https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json',
    'facebook/xmod-base-60-125k': 'https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json',
    'facebook/xmod-base-60-265k': 'https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json',
    'facebook/xmod-base-75-125k': 'https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json',
    'facebook/xmod-base-75-269k': 'https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json',
}







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			_UpperCamelCase		:       List[str] 			=		'xmod'




			def __init__(				self   ,    snake_case=30522   ,    snake_case=768   ,    snake_case=12   ,    snake_case=12   ,    snake_case=3072   ,    snake_case="gelu"   ,    snake_case=0.1   ,    snake_case=0.1   ,    snake_case=512   ,    snake_case=2   ,    snake_case=0.02   ,    snake_case=1E-1_2   ,    snake_case=1   ,    snake_case=0   ,    snake_case=2   ,    snake_case="absolute"   ,    snake_case=True   ,    snake_case=None   ,    snake_case=False   ,    snake_case=2   ,    snake_case=False   ,    snake_case=True   ,    snake_case=True   ,    snake_case=("en_XX",)   ,    snake_case=None   ,    **snake_case   ,    ):

						'''simple docstring'''


						super().__init__(pad_token_id=snake_case   ,    bos_token_id=snake_case   ,    eos_token_id=snake_case   ,    **snake_case  )

						UpperCamelCase__       =							vocab_size
						UpperCamelCase__       =							hidden_size
						UpperCamelCase__       =							num_hidden_layers
						UpperCamelCase__       =							num_attention_heads
						UpperCamelCase__       =							hidden_act
						UpperCamelCase__       =							intermediate_size
						UpperCamelCase__       =							hidden_dropout_prob
						UpperCamelCase__       =							attention_probs_dropout_prob
						UpperCamelCase__       =							max_position_embeddings
						UpperCamelCase__       =							type_vocab_size
						UpperCamelCase__       =							initializer_range
						UpperCamelCase__       =							layer_norm_eps
						UpperCamelCase__       =							position_embedding_type
						UpperCamelCase__       =							use_cache
						UpperCamelCase__       =							classifier_dropout
						UpperCamelCase__       =							pre_norm
						UpperCamelCase__       =							adapter_reduction_factor
						UpperCamelCase__       =							adapter_layer_norm
						UpperCamelCase__       =							adapter_reuse_layer_norm
						UpperCamelCase__       =							ln_before_adapter
						UpperCamelCase__       =							list(snake_case  )
						UpperCamelCase__       =							default_language







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			@property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						if self.task == "multiple-choice":
									UpperCamelCase__       =							{0: "batch", 1: "choice", 2: "sequence"}
						else:
									UpperCamelCase__       =							{0: "batch", 1: "sequence"}
						return OrderedDict(
						    [
						        ("input_ids", dynamic_axis),
						        ("attention_mask", dynamic_axis),
						    ]  )


