





from ..utils import is_flax_available, is_torch_available


if is_torch_available():
					from .autoencoder_kl import AutoencoderKL
					from .controlnet import ControlNetModel
					from .dual_transformer_ad import DualTransformeraDModel
					from .modeling_utils import ModelMixin
					from .prior_transformer import PriorTransformer
					from .ta_film_transformer import TaFilmDecoder
					from .transformer_ad import TransformeraDModel
					from .unet_ad import UNetaDModel
					from .unet_ad import UNetaDModel
					from .unet_ad_condition import UNetaDConditionModel
					from .unet_ad_condition import UNetaDConditionModel
					from .vq_model import VQModel

if is_flax_available():
					from .controlnet_flax import FlaxControlNetModel
					from .unet_ad_condition_flax import FlaxUNetaDConditionModel
					from .vae_flax import FlaxAutoencoderKL











import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import DetrConfig, MaskFormerConfig, SwinConfig, 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 MaskFormerForInstanceSegmentation, MaskFormerModel

					if is_vision_available():
										from transformers import MaskFormerImageProcessor

if is_vision_available():
					from PIL import Image

class        __SCREAMING_SNAKE_CASE					:






							def __init__(						self  , __lowerCAmelCase  , __lowerCAmelCase=2  , __lowerCAmelCase=True  , __lowerCAmelCase=False  , __lowerCAmelCase=10  , __lowerCAmelCase=3  , __lowerCAmelCase=32 * 4  , __lowerCAmelCase=32 * 6  , __lowerCAmelCase=4  , __lowerCAmelCase=32  , ):
														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__				     =			mask_feature_size






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__				     =			floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]  ).to(
														    __lowerCAmelCase  )

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

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

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






							def 				_lowerCamelCase		(						self  ):
														return MaskFormerConfig.from_backbone_and_decoder_configs(
														    backbone_config=SwinConfig(
														        depths=[1, 1, 1, 1]  , )  , decoder_config=DetrConfig(
														        decoder_ffn_dim=128  , num_queries=self.num_queries  , decoder_attention_heads=2  , d_model=self.mask_feature_size  , )  , mask_feature_size=self.mask_feature_size  , fpn_feature_size=self.mask_feature_size  , num_channels=self.num_channels  , num_labels=self.num_labels  , )






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






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase  ):
														UpperCamelCase__				     =			output.encoder_hidden_states
														UpperCamelCase__				     =			output.pixel_decoder_hidden_states
														UpperCamelCase__				     =			output.transformer_decoder_hidden_states

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






							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase=False  ):
														with torch.no_grad():
																					UpperCamelCase__				     =			MaskFormerModel(config=__lowerCAmelCase  )
																					model.to(__lowerCAmelCase  )
																					model.eval()

																					UpperCamelCase__				     =			model(pixel_values=__lowerCAmelCase  , pixel_mask=__lowerCAmelCase  )
																					UpperCamelCase__				     =			model(__lowerCAmelCase  , output_hidden_states=__lowerCAmelCase  )
														# the correct shape of output.transformer_decoder_hidden_states ensure the correcteness of the
														# encoder and pixel decoder
														self.parent.assertEqual(
														    output.transformer_decoder_last_hidden_state.shape  , (self.batch_size, self.num_queries, self.mask_feature_size)  , )
														# 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(__lowerCAmelCase  , __lowerCAmelCase  )







							def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase  , __lowerCAmelCase  ):
														UpperCamelCase__				     =			MaskFormerForInstanceSegmentation(config=__lowerCAmelCase  )
														model.to(__lowerCAmelCase  )
														model.eval()

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

																					comm_check_on_output(__lowerCAmelCase  )

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

														comm_check_on_output(__lowerCAmelCase  )

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






@require_torch
class        __SCREAMING_SNAKE_CASE					(     _a   ,		_a   ,		unittest.TestCase       ):
							snake_case     :      str						     =       (MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
							snake_case     :      str						     =       (
							    {"""feature-extraction""": MaskFormerModel, """image-segmentation""": MaskFormerForInstanceSegmentation}
							    if is_torch_available()
							    else {}
							)

							snake_case     :      Tuple						     =       False
							snake_case     :      str						     =       False
							snake_case     :      Any						     =       False
							snake_case     :      Union[str, Any]						     =       False






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__				     =			MaskFormerModelTester(self  )
														UpperCamelCase__				     =			ConfigTester(self  , config_class=__lowerCAmelCase  , has_text_modality=__lowerCAmelCase  )






							def 				_lowerCamelCase		(						self  ):
														self.config_tester.run_common_tests()






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__     ,     UpperCamelCase__				     =			self.model_tester.prepare_config_and_inputs_for_common()
														self.model_tester.create_and_check_maskformer_model(__lowerCAmelCase  , **__lowerCAmelCase  , output_hidden_states=__lowerCAmelCase  )






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__				     =			self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*__lowerCAmelCase  )






							@unittest.skip(reason="""MaskFormer does not use inputs_embeds"""  )
							def 				_lowerCamelCase		(						self  ):
														pass






							@unittest.skip(reason="""MaskFormer does not have a get_input_embeddings method"""  )
							def 				_lowerCamelCase		(						self  ):
														pass






							@unittest.skip(reason="""MaskFormer is not a generative model"""  )
							def 				_lowerCamelCase		(						self  ):
														pass






							@unittest.skip(reason="""MaskFormer does not use token embeddings"""  )
							def 				_lowerCamelCase		(						self  ):
														pass






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






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






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__     ,     UpperCamelCase__				     =			self.model_tester.prepare_config_and_inputs_for_common()

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






							@slow
							def 				_lowerCamelCase		(						self  ):
														for model_name in ["facebook/maskformer-swin-small-coco"]:
																					UpperCamelCase__				     =			MaskFormerModel.from_pretrained(__lowerCAmelCase  )
																					self.assertIsNotNone(__lowerCAmelCase  )






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__				     =			(self.model_tester.min_size,) * 2
														UpperCamelCase__				     =			{
														    """pixel_values""": torch.randn((2, 3, *size)  , device=__lowerCAmelCase  ),
														    """mask_labels""": torch.randn((2, 10, *size)  , device=__lowerCAmelCase  ),
														    """class_labels""": torch.zeros(2  , 10  , device=__lowerCAmelCase  ).long(),
														}

														UpperCamelCase__				     =			MaskFormerForInstanceSegmentation(MaskFormerConfig()  ).to(__lowerCAmelCase  )
														UpperCamelCase__				     =			model(**__lowerCAmelCase  )
														self.assertTrue(outputs.loss is not None  )






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__     ,     UpperCamelCase__				     =			self.model_tester.prepare_config_and_inputs_for_common()
														self.model_tester.create_and_check_maskformer_model(__lowerCAmelCase  , **__lowerCAmelCase  , output_hidden_states=__lowerCAmelCase  )






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__     ,     UpperCamelCase__				     =			self.model_tester.prepare_config_and_inputs_for_common()

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






							def 				_lowerCamelCase		(						self  ):
														if not self.model_tester.is_training:
																					return
														# only MaskFormerForInstanceSegmentation has the loss
														UpperCamelCase__				     =			self.all_model_classes[1]
														UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__				     =			self.model_tester.prepare_config_and_inputs()

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

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







							def 				_lowerCamelCase		(						self  ):
														# only MaskFormerForInstanceSegmentation has the loss
														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(__lowerCAmelCase  )
														model.to(__lowerCAmelCase  )
														model.train()

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

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

														UpperCamelCase__				     =			outputs.pixel_decoder_hidden_states[0]
														pixel_decoder_hidden_states.retain_grad()
														# we requires_grad=True in inputs_embeds (line 2152), the original implementation don't
														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=__lowerCAmelCase  )

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






							"""simple docstring"""





							UpperCamelCase__				     =			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""       )
							return image

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






							@cached_property
							def 				_lowerCamelCase		(						self  ):
														return (
														    MaskFormerImageProcessor.from_pretrained("""facebook/maskformer-swin-small-coco"""  )
														    if is_vision_available()
														    else None
														)






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__				     =			MaskFormerModel.from_pretrained("""facebook/maskformer-swin-small-coco"""  ).to(__lowerCAmelCase  )
														UpperCamelCase__				     =			self.default_image_processor
														UpperCamelCase__				     =			prepare_img()
														UpperCamelCase__				     =			image_processor(__lowerCAmelCase  , return_tensors="""pt"""  ).to(__lowerCAmelCase  )
														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(__lowerCAmelCase  , (1, 3, 800, 1088)  )

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

														UpperCamelCase__				     =			torch.tensor(
														    [[-0.0482, 0.9228, 0.4951], [-0.2547, 0.8017, 0.8527], [-0.0069, 0.3385, -0.0089]]  ).to(__lowerCAmelCase  )
														self.assertTrue(
														    torch.allclose(
														        outputs.encoder_last_hidden_state[0, 0, :3, :3]  , __lowerCAmelCase  , atol=__lowerCAmelCase  )  )

														UpperCamelCase__				     =			torch.tensor(
														    [[-0.8422, -0.8434, -0.9718], [-1.0144, -0.5565, -0.4195], [-1.0038, -0.4484, -0.1961]]  ).to(__lowerCAmelCase  )
														self.assertTrue(
														    torch.allclose(
														        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]  , __lowerCAmelCase  , atol=__lowerCAmelCase  )  )

														UpperCamelCase__				     =			torch.tensor(
														    [[0.2852, -0.0159, 0.9735], [0.6254, 0.1858, 0.8529], [-0.0680, -0.4116, 1.8413]]  ).to(__lowerCAmelCase  )
														self.assertTrue(
														    torch.allclose(
														        outputs.transformer_decoder_last_hidden_state[0, :3, :3]  , __lowerCAmelCase  , atol=__lowerCAmelCase  )  )






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__				     =			(
														    MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-swin-small-coco"""  )
														    .to(__lowerCAmelCase  )
														    .eval()
														)
														UpperCamelCase__				     =			self.default_image_processor
														UpperCamelCase__				     =			prepare_img()
														UpperCamelCase__				     =			image_processor(__lowerCAmelCase  , return_tensors="""pt"""  ).to(__lowerCAmelCase  )
														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(__lowerCAmelCase  , (1, 3, 800, 1088)  )

														with torch.no_grad():
																					UpperCamelCase__				     =			model(**__lowerCAmelCase  )
														# masks_queries_logits
														UpperCamelCase__				     =			outputs.masks_queries_logits
														self.assertEqual(
														    masks_queries_logits.shape  , (1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)  , )
														UpperCamelCase__				     =			[
														    [-1.373_7124, -1.772_4937, -1.936_4233],
														    [-1.597_7281, -1.986_7939, -2.152_3695],
														    [-1.579_5398, -1.926_9832, -2.09_3942],
														]
														UpperCamelCase__				     =			torch.tensor(__lowerCAmelCase  ).to(__lowerCAmelCase  )
														self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]  , __lowerCAmelCase  , atol=__lowerCAmelCase  )  )
														# class_queries_logits
														UpperCamelCase__				     =			outputs.class_queries_logits
														self.assertEqual(
														    class_queries_logits.shape  , (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)  )
														UpperCamelCase__				     =			torch.tensor(
														    [
														        [1.6512E00, -5.2572E00, -3.3519E00],
														        [3.6169E-02, -5.9025E00, -2.9313E00],
														        [1.0766E-04, -7.7630E00, -5.1263E00],
														    ]  ).to(__lowerCAmelCase  )
														self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]  , __lowerCAmelCase  , atol=__lowerCAmelCase  )  )






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__				     =			(
														    MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-resnet101-coco-stuff"""  )
														    .to(__lowerCAmelCase  )
														    .eval()
														)
														UpperCamelCase__				     =			self.default_image_processor
														UpperCamelCase__				     =			prepare_img()
														UpperCamelCase__				     =			image_processor(__lowerCAmelCase  , return_tensors="""pt"""  ).to(__lowerCAmelCase  )
														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(__lowerCAmelCase  , (1, 3, 800, 1088)  )

														with torch.no_grad():
																					UpperCamelCase__				     =			model(**__lowerCAmelCase  )
														# masks_queries_logits
														UpperCamelCase__				     =			outputs.masks_queries_logits
														self.assertEqual(
														    masks_queries_logits.shape  , (1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)  , )
														UpperCamelCase__				     =			[[-0.9046, -2.6366, -4.6062], [-3.4179, -5.7890, -8.8057], [-4.9179, -7.6560, -10.7711]]
														UpperCamelCase__				     =			torch.tensor(__lowerCAmelCase  ).to(__lowerCAmelCase  )
														self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]  , __lowerCAmelCase  , atol=__lowerCAmelCase  )  )
														# class_queries_logits
														UpperCamelCase__				     =			outputs.class_queries_logits
														self.assertEqual(
														    class_queries_logits.shape  , (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)  )
														UpperCamelCase__				     =			torch.tensor(
														    [[4.7188, -3.2585, -2.8857], [6.6871, -2.9181, -1.2487], [7.2449, -2.2764, -2.1874]]  ).to(__lowerCAmelCase  )
														self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]  , __lowerCAmelCase  , atol=__lowerCAmelCase  )  )







							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__				     =			(
														    MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-swin-small-coco"""  )
														    .to(__lowerCAmelCase  )
														    .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(__lowerCAmelCase  )
														UpperCamelCase__				     =			[el.to(__lowerCAmelCase  ) for el in inputs["""mask_labels"""]]
														UpperCamelCase__				     =			[el.to(__lowerCAmelCase  ) for el in inputs["""class_labels"""]]

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

														self.assertTrue(outputs.loss is not None  )




