






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       __lowercase      :







     '''simple docstring'''




     def __init__(				self       :			Any ,	UpperCamelCase_       :			Optional[Any] ,	UpperCamelCase_       :			Tuple=2 ,	UpperCamelCase_       :			Tuple=True ,	UpperCamelCase_       :			Any=False ,	UpperCamelCase_       :			Any=10 ,	UpperCamelCase_       :			List[Any]=3 ,	UpperCamelCase_       :			str=32 * 8 ,	UpperCamelCase_       :			List[str]=32 * 8 ,	UpperCamelCase_       :			int=4 ,	UpperCamelCase_       :			List[str]=64 ,	):


         """simple docstring"""



         __A         =	parent
         __A         =	batch_size
         __A         =	is_training
         __A         =	use_auxiliary_loss
         __A         =	num_queries
         __A         =	num_channels
         __A         =	min_size
         __A         =	max_size
         __A         =	num_labels
         __A         =	hidden_dim
         __A         =	hidden_dim




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]				).to(
             UpperCamelCase_				)

         __A         =	torch.ones([self.batch_size, self.min_size, self.max_size] ,	device=UpperCamelCase_				)

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

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




     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         __A         =	MaskaFormerConfig(
             hidden_size=self.hidden_dim ,	)
         __A         =	self.num_queries
         __A         =	self.num_labels

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

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




     def 		lowerCAmelCase_     (				self       :			Tuple				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			Tuple ,	UpperCamelCase_       :			Dict ,	UpperCamelCase_       :			int				):


         """simple docstring"""



         __A         =	output.encoder_hidden_states
         __A         =	output.pixel_decoder_hidden_states
         __A         =	output.transformer_decoder_hidden_states

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




     def 		lowerCAmelCase_     (				self       :			Any ,	UpperCamelCase_       :			Optional[int] ,	UpperCamelCase_       :			Tuple ,	UpperCamelCase_       :			str ,	UpperCamelCase_       :			Any=False				):


         """simple docstring"""



         with torch.no_grad():
             __A         =	MaskaFormerModel(config=UpperCamelCase_				)
             model.to(UpperCamelCase_				)
             model.eval()

             __A         =	model(pixel_values=UpperCamelCase_ ,	pixel_mask=UpperCamelCase_				)
             __A         =	model(UpperCamelCase_ ,	output_hidden_states=UpperCamelCase_				)

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





     def 		lowerCAmelCase_     (				self       :			str ,	UpperCamelCase_       :			Any ,	UpperCamelCase_       :			Union[str, Any] ,	UpperCamelCase_       :			Union[str, Any] ,	UpperCamelCase_       :			List[Any] ,	UpperCamelCase_       :			Dict				):


         """simple docstring"""



         __A         =	MaskaFormerForUniversalSegmentation(config=UpperCamelCase_				)
         model.to(UpperCamelCase_				)
         model.eval()

         def comm_check_on_output(UpperCamelCase_       :			List[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():
             __A         =	model(pixel_values=UpperCamelCase_ ,	pixel_mask=UpperCamelCase_				)
             __A         =	model(UpperCamelCase_				)

             comm_check_on_output(UpperCamelCase_				)

             __A         =	model(
                 pixel_values=UpperCamelCase_ ,	pixel_mask=UpperCamelCase_ ,	mask_labels=UpperCamelCase_ ,	class_labels=UpperCamelCase_				)

         comm_check_on_output(UpperCamelCase_				)

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


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







     '''simple docstring'''
     SCREAMING_SNAKE_CASE   						=				(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
     SCREAMING_SNAKE_CASE   						=				{"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False
     SCREAMING_SNAKE_CASE   						=				False




     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         __A         =	MaskaFormerModelTester(self				)
         __A         =	ConfigTester(self ,	config_class=UpperCamelCase_ ,	has_text_modality=UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         self.config_tester.run_common_tests()




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()
         self.model_tester.create_and_check_maskaformer_model(UpperCamelCase_ ,	**UpperCamelCase_ ,	output_hidden_states=UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*UpperCamelCase_				)




     @unittest.skip(reason="""Mask2Former does not use inputs_embeds"""				)
     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         pass




     @unittest.skip(reason="""Mask2Former does not have a get_input_embeddings method"""				)
     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         pass




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


         """simple docstring"""



         pass




     @unittest.skip(reason="""Mask2Former does not use token embeddings"""				)
     def 		lowerCAmelCase_     (				self       :			Union[str, Any]				):


         """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 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         pass




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


         """simple docstring"""



         pass




     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()

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

             __A         =	["""pixel_values"""]
             self.assertListEqual(arg_names[:1] ,	UpperCamelCase_				)




     @slow
     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			int				):


         """simple docstring"""



         __A         =	(self.model_tester.min_size,) * 2
         __A         =	{
             """pixel_values""": torch.randn((2, 3, *size) ,	device=UpperCamelCase_				),
             """mask_labels""": torch.randn((2, 10, *size) ,	device=UpperCamelCase_				),
             """class_labels""": torch.zeros(2 ,	10 ,	device=UpperCamelCase_				).long(),
         }
         __A         =	self.model_tester.get_config()

         __A         =	MaskaFormerForUniversalSegmentation(UpperCamelCase_				).to(UpperCamelCase_				)
         __A         =	model(**UpperCamelCase_				)
         self.assertTrue(outputs.loss is not None				)




     def 		lowerCAmelCase_     (				self       :			Union[str, Any]				):


         """simple docstring"""



         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()
         self.model_tester.create_and_check_maskaformer_model(UpperCamelCase_ ,	**UpperCamelCase_ ,	output_hidden_states=UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A    ,	__A         =	self.model_tester.prepare_config_and_inputs_for_common()

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




     def 		lowerCAmelCase_     (				self       :			Tuple				):


         """simple docstring"""



         if not self.model_tester.is_training:
             return

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

         __A         =	model_class(UpperCamelCase_				)
         model.to(UpperCamelCase_				)
         model.train()

         __A         =	model(UpperCamelCase_ ,	mask_labels=UpperCamelCase_ ,	class_labels=UpperCamelCase_				).loss
         loss.backward()





     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



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

         __A         =	model_class(UpperCamelCase_				).to(UpperCamelCase_				)
         model.train()

         __A         =	model(UpperCamelCase_ ,	mask_labels=UpperCamelCase_ ,	class_labels=UpperCamelCase_				)

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

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

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

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

         outputs.loss.backward(retain_graph=UpperCamelCase_				)

         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      :       Any				  =   1e-4



def 			_SCREAMING_SNAKE_CASE						( )						->      Union[str, Any]:





    """simple docstring"""

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






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







     '''simple docstring'''




     @cached_property
     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



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




     @cached_property
     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



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




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



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

         with torch.no_grad():
             __A         =	model(**UpperCamelCase_				)

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

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

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




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



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

         with torch.no_grad():
             __A         =	model(**UpperCamelCase_				)
         # masks_queries_logits
         __A         =	outputs.masks_queries_logits
         self.assertEqual(
             masks_queries_logits.shape ,	(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)				)
         __A         =	[
             [-8.7839, -9.0056, -8.8121],
             [-7.4104, -7.0313, -6.5401],
             [-6.6105, -6.3427, -6.4675],
         ]
         __A         =	torch.tensor(UpperCamelCase_				).to(UpperCamelCase_				)
         self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3] ,	UpperCamelCase_ ,	atol=UpperCamelCase_				)				)
         # class_queries_logits
         __A         =	outputs.class_queries_logits
         self.assertEqual(class_queries_logits.shape ,	(1, model.config.num_queries, model.config.num_labels + 1)				)
         __A         =	torch.tensor(
             [
                 [1.8324, -8.0835, -4.1922],
                 [0.8450, -9.0050, -3.6053],
                 [0.3045, -7.7293, -3.0275],
             ]				).to(UpperCamelCase_				)
         self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3] ,	UpperCamelCase_ ,	atol=UpperCamelCase_				)				)





     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         __A         =	MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints				).to(UpperCamelCase_				).eval()
         __A         =	self.default_image_processor

         __A         =	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""" ,	)

         __A         =	inputs["""pixel_values"""].to(UpperCamelCase_				)
         __A         =	[el.to(UpperCamelCase_				) for el in inputs["""mask_labels"""]]
         __A         =	[el.to(UpperCamelCase_				) for el in inputs["""class_labels"""]]

         with torch.no_grad():
             __A         =	model(**UpperCamelCase_				)

         self.assertTrue(outputs.loss is not None				)












import platform
from argparse import ArgumentParser

import huggingface_hub

from .. import __version__ as version
from ..utils import is_accelerate_available, is_torch_available, is_transformers_available, is_xformers_available
from . import BaseDiffusersCLICommand



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Any    )						->      Optional[int]:





    """simple docstring"""

    return EnvironmentCommand()






class       __lowercase      (	lowercase_  ):







     '''simple docstring'''




     @staticmethod
     def 		lowerCAmelCase_     (				UpperCamelCase_       :			ArgumentParser				):


         """simple docstring"""



         __A         =	parser.add_parser("""env"""				)
         download_parser.set_defaults(func=UpperCamelCase_				)




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	huggingface_hub.__version__

         __A         =	"""not installed"""
         __A         =	"""NA"""
         if is_torch_available():
             import torch

             __A         =	torch.__version__
             __A         =	torch.cuda.is_available()

         __A         =	"""not installed"""
         if is_transformers_available():
             import transformers

             __A         =	transformers.__version__

         __A         =	"""not installed"""
         if is_accelerate_available():
             import accelerate

             __A         =	accelerate.__version__

         __A         =	"""not installed"""
         if is_xformers_available():
             import xformers

             __A         =	xformers.__version__

         __A         =	{
             """`diffusers` version""": version,
             """Platform""": platform.platform(),
             """Python version""": platform.python_version(),
             """PyTorch version (GPU?)""": F"{pt_version} ({pt_cuda_available})",
             """Huggingface_hub version""": hub_version,
             """Transformers version""": transformers_version,
             """Accelerate version""": accelerate_version,
             """xFormers version""": xformers_version,
             """Using GPU in script?""": """<fill in>""",
             """Using distributed or parallel set-up in script?""": """<fill in>""",
         }

         print("""\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n"""				)
         print(self.format_dict(UpperCamelCase_				)				)

         return info





     @staticmethod
     def 		lowerCAmelCase_     (				UpperCamelCase_       :			str				):


         """simple docstring"""



         return "\n".join([F"- {prop}: {val}" for prop, val in d.items()]				) + "\n"




