

'''simple docstring'''





import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

						if is_vision_available():
												from transformers import MaskaFormerImageProcessor

if is_vision_available():
						from PIL import Image

class 				__UpperCamelCase				:







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

						'''simple docstring'''





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







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





						__magic_name__							=    floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]      ).to(
						    _lowercase      )

						__magic_name__							=    torch.ones([self.batch_size, self.min_size, self.max_size]  , device=_lowercase      )

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

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







	def 	_SCREAMING_SNAKE_CASE							(   self: List[str]      ):

						'''simple docstring'''





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

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

						__magic_name__							=    64
						__magic_name__							=    1_28
						__magic_name__							=    self.hidden_dim
						__magic_name__							=    self.hidden_dim
						__magic_name__							=    self.hidden_dim
						return config







	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[int]      ):

						'''simple docstring'''





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







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple  , __UpperCamelCase: List[str]  , __UpperCamelCase: Dict      ):

						'''simple docstring'''





						__magic_name__							=    output.encoder_hidden_states
						__magic_name__							=    output.pixel_decoder_hidden_states
						__magic_name__							=    output.transformer_decoder_hidden_states

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







	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[Any]  , __UpperCamelCase: Dict  , __UpperCamelCase: int  , __UpperCamelCase: Any  , __UpperCamelCase: Optional[int]=False      ):

						'''simple docstring'''





						with torch.no_grad():
											__magic_name__							=    MaskaFormerModel(config=_lowercase      )
											model.to(_lowercase      )
											model.eval()

											__magic_name__							=    model(pixel_values=_lowercase  , pixel_mask=_lowercase      )
											__magic_name__							=    model(_lowercase  , output_hidden_states=_lowercase      )

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







	def 	_SCREAMING_SNAKE_CASE							(   self: Any  , __UpperCamelCase: int  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: List[str]  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Optional[Any]      ):

						'''simple docstring'''





						__magic_name__							=    MaskaFormerForUniversalSegmentation(config=_lowercase      )
						model.to(_lowercase      )
						model.eval()

						def comm_check_on_output(__UpperCamelCase: int      ):
											# 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():
											__magic_name__							=    model(pixel_values=_lowercase  , pixel_mask=_lowercase      )
											__magic_name__							=    model(_lowercase      )

											comm_check_on_output(_lowercase      )

											__magic_name__							=    model(
											    pixel_values=_lowercase  , pixel_mask=_lowercase  , mask_labels=_lowercase  , class_labels=_lowercase      )

						comm_check_on_output(_lowercase      )

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




@require_torch
class 				__UpperCamelCase				(					SCREAMING_SNAKE_CASE__  ,						SCREAMING_SNAKE_CASE__  ,						unittest.TestCase			):
	_lowercase  :   Optional[int]       				=  (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
	_lowercase  :   Union[str, Any]       				=  {"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

	_lowercase  :   List[str]       				=  False
	_lowercase  :   Dict       				=  False
	_lowercase  :   str       				=  False
	_lowercase  :   str       				=  False







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

						'''simple docstring'''





						__magic_name__							=    MaskaFormerModelTester(self      )
						__magic_name__							=    ConfigTester(self  , config_class=_lowercase  , has_text_modality=_lowercase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						self.config_tester.run_common_tests()







	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[Any]      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs_for_common()
						self.model_tester.create_and_check_maskaformer_model(_lowercase  , **_lowercase  , output_hidden_states=_lowercase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: str      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*_lowercase      )







	@unittest.skip(reason='Mask2Former does not use inputs_embeds'      )
	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





						pass







	@unittest.skip(reason='Mask2Former does not have a get_input_embeddings method'      )
	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						pass







	@unittest.skip(reason='Mask2Former is not a generative model'      )
	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						pass







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

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

						'''simple docstring'''





						pass







	@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'      )
	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						pass







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs_for_common()

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

											__magic_name__							=    ["""pixel_values"""]
											self.assertListEqual(arg_names[:1]  , _lowercase      )







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





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







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





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

						__magic_name__							=    MaskaFormerForUniversalSegmentation(_lowercase      ).to(_lowercase      )
						__magic_name__							=    model(**_lowercase      )
						self.assertTrue(outputs.loss is not None      )







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

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs_for_common()
						self.model_tester.create_and_check_maskaformer_model(_lowercase  , **_lowercase  , output_hidden_states=_lowercase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs_for_common()

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







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





						if not self.model_tester.is_training:
											return

						__magic_name__							=    self.all_model_classes[1]
						__magic_name__							=    self.model_tester.prepare_config_and_inputs()

						__magic_name__							=    model_class(_lowercase      )
						model.to(_lowercase      )
						model.train()

						__magic_name__							=    model(_lowercase  , mask_labels=_lowercase  , class_labels=_lowercase      ).loss
						loss.backward()







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





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

						__magic_name__							=    model_class(_lowercase      ).to(_lowercase      )
						model.train()

						__magic_name__							=    model(_lowercase  , mask_labels=_lowercase  , class_labels=_lowercase      )

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

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

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

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

						outputs.loss.backward(retain_graph=_lowercase      )

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




					'''simple docstring'''







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


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







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

						'''simple docstring'''





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







	@cached_property
	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]      ):

						'''simple docstring'''





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







	def 	_SCREAMING_SNAKE_CASE							(   self: List[str]      ):

						'''simple docstring'''





						__magic_name__							=    MaskaFormerModel.from_pretrained(self.model_checkpoints      ).to(_lowercase      )
						__magic_name__							=    self.default_image_processor
						__magic_name__							=    prepare_img()
						__magic_name__							=    image_processor(_lowercase  , return_tensors='pt'      ).to(_lowercase      )
						__magic_name__							=    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(_lowercase  , (1, 3, 3_84, 3_84)      )

						with torch.no_grad():
											__magic_name__							=    model(**_lowercase      )

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

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

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







	def 	_SCREAMING_SNAKE_CASE							(   self: List[str]      ):

						'''simple docstring'''





						__magic_name__							=    MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints      ).to(_lowercase      ).eval()
						__magic_name__							=    self.default_image_processor
						__magic_name__							=    prepare_img()
						__magic_name__							=    image_processor(_lowercase  , return_tensors='pt'      ).to(_lowercase      )
						__magic_name__							=    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(_lowercase  , (1, 3, 3_84, 3_84)      )

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







	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[Any]      ):

						'''simple docstring'''





						__magic_name__							=    MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints      ).to(_lowercase      ).eval()
						__magic_name__							=    self.default_image_processor

						__magic_name__							=    image_processor(
						    [np.zeros((3, 8_00, 13_33)      ), np.zeros((3, 8_00, 13_33)      )]  , segmentation_maps=[np.zeros((3_84, 3_84)      ).astype(np.floataa      ), np.zeros((3_84, 3_84)      ).astype(np.floataa      )]  , return_tensors='pt'  , )

						__magic_name__							=    inputs["""pixel_values"""].to(_lowercase      )
						__magic_name__							=    [el.to(_lowercase      ) for el in inputs["""mask_labels"""]]
						__magic_name__							=    [el.to(_lowercase      ) for el in inputs["""class_labels"""]]

						with torch.no_grad():
											__magic_name__							=    model(**_lowercase      )

						self.assertTrue(outputs.loss is not None      )







import inspect
import os
import unittest
from dataclasses import dataclass

import torch

from accelerate import Accelerator, DistributedDataParallelKwargs, GradScalerKwargs
from accelerate.state import AcceleratorState
from accelerate.test_utils import execute_subprocess_async, require_cuda, require_multi_gpu
from accelerate.utils import KwargsHandler

@dataclass
class 				__UpperCamelCase				(					SCREAMING_SNAKE_CASE			):
	_lowercase  :   int       				=  0
	_lowercase  :   bool       				=  False
	_lowercase  :   float       				=  3.0




class 				__UpperCamelCase				(					unittest.TestCase			):







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						self.assertDictEqual(MockClass().to_kwargs()  , {}      )
						self.assertDictEqual(MockClass(a=2      ).to_kwargs()  , {'a': 2}      )
						self.assertDictEqual(MockClass(a=2  , b=__UpperCamelCase      ).to_kwargs()  , {'a': 2, 'b': True}      )
						self.assertDictEqual(MockClass(a=2  , c=2.25      ).to_kwargs()  , {'a': 2, 'c': 2.25}      )







	@require_cuda
	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





						__magic_name__							=    GradScalerKwargs(init_scale=10_24  , growth_factor=2      )
						AcceleratorState._reset_state()
						__magic_name__							=    Accelerator(mixed_precision='fp16'  , kwargs_handlers=[scaler_handler]      )
						print(accelerator.use_fpaa      )
						__magic_name__							=    accelerator.scaler

						# Check the kwargs have been applied
						self.assertEqual(scaler._init_scale  , 1024.0      )
						self.assertEqual(scaler._growth_factor  , 2.0      )

						# Check the other values are at the default
						self.assertEqual(scaler._backoff_factor  , 0.5      )
						self.assertEqual(scaler._growth_interval  , 20_00      )
						self.assertEqual(scaler._enabled  , __UpperCamelCase      )







	@require_multi_gpu
	def 	_SCREAMING_SNAKE_CASE							(   self: Optional[Any]      ):

						'''simple docstring'''





						__magic_name__							=    ['torchrun', F'--nproc_per_node={torch.cuda.device_count()}', inspect.getfile(self.__class__      )]
						execute_subprocess_async(__UpperCamelCase  , env=os.environ.copy()      )


if __name__ == "__main__":
						A__						=							DistributedDataParallelKwargs(bucket_cap_mb=15, find_unused_parameters=True)
						A__						=							Accelerator(kwargs_handlers=[ddp_scaler])
						A__						=							torch.nn.Linear(100, 200)
						A__						=							accelerator.prepare(model)

						# Check the values changed in kwargs
						A__						=							""
						A__						=							model.bucket_bytes_cap // (1024 * 1024)
						if observed_bucket_cap_map != 15:
												error_msg += f"Kwargs badly passed, should have `15` but found {observed_bucket_cap_map}.\n"
						if model.find_unused_parameters is not True:
												error_msg += f"Kwargs badly passed, should have `True` but found {model.find_unused_parameters}.\n"

						# Check the values of the defaults
						if model.dim != 0:
												error_msg += f"Default value not respected, should have `0` but found {model.dim}.\n"
						if model.broadcast_buffers is not True:
												error_msg += f"Default value not respected, should have `True` but found {model.broadcast_buffers}.\n"
						if model.gradient_as_bucket_view is not False:
												error_msg += f"Default value not respected, should have `False` but found {model.gradient_as_bucket_view}.\n"

						# Raise error at the end to make sure we don't stop at the first failure.
						if len(error_msg) > 0:
												raise ValueError(error_msg)




