




'''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 					_a :





     '''simple docstring'''



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







     def 			snake_case_      (				self							)      ->				List[Any]:
          snake_case			:      Union[str, Any]        =			floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]							).to(
              __a							)

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

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

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







     def 			snake_case_      (				self							)      ->				Tuple:
          snake_case			:      Union[str, Any]        =			MaskaFormerConfig(
              hidden_size=self.hidden_dim							,)
          snake_case			:      Dict        =			self.num_queries
          snake_case			:      Any        =			self.num_labels

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

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







     def 			snake_case_      (				self							)      ->				Optional[int]:
          snake_case ,		snake_case ,		snake_case ,		snake_case ,		snake_case			:      int        =			self.prepare_config_and_inputs()
          snake_case			:      List[str]        =			{"""pixel_values""": pixel_values, """pixel_mask""": pixel_mask}
          return config, inputs_dict







     def 			snake_case_      (				self							,__a							,__a							)      ->				Union[str, Any]:
          snake_case			:      Any        =			output.encoder_hidden_states
          snake_case			:      List[str]        =			output.pixel_decoder_hidden_states
          snake_case			:      Tuple        =			output.transformer_decoder_hidden_states

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







     def 			snake_case_      (				self							,__a							,__a							,__a							,__a=False							)      ->				Optional[Any]:
          with torch.no_grad():
               snake_case			:      Any        =			MaskaFormerModel(config=__a							)
               model.to(__a							)
               model.eval()

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

          self.parent.assertEqual(
              output.transformer_decoder_last_hidden_state.shape							,(self.batch_size, self.num_queries, self.hidden_dim)							,)
          # let's ensure the other two hidden state exists
          self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None							)
          self.parent.assertTrue(output.encoder_last_hidden_state is not None							)

          if output_hidden_states:
               self.check_output_hidden_state(__a							,__a							)







     def 			snake_case_      (				self							,__a							,__a							,__a							,__a							,__a							)      ->				Union[str, Any]:
          snake_case			:      Optional[Any]        =			MaskaFormerForUniversalSegmentation(config=__a							)
          model.to(__a							)
          model.eval()

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

               comm_check_on_output(__a							)

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

          comm_check_on_output(__a							)

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







@require_torch
class 					_a (a__,		a__,		unittest.TestCase	):





     '''simple docstring'''



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

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







     def 			snake_case_      (				self							)      ->				Union[str, Any]:
          snake_case			:      Optional[Any]        =			MaskaFormerModelTester(self							)
          snake_case			:      List[Any]        =			ConfigTester(self							,config_class=__a							,has_text_modality=__a							)







     def 			snake_case_      (				self							)      ->				Any:
          self.config_tester.run_common_tests()







     def 			snake_case_      (				self							)      ->				List[Any]:
          snake_case ,		snake_case			:      int        =			self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.create_and_check_maskaformer_model(__a							,**__a							,output_hidden_states=__a							)







     def 			snake_case_      (				self							)      ->				int:
          snake_case			:      Any        =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*__a							)







     @unittest.skip(reason="""Mask2Former does not use inputs_embeds"""							)
     def 			snake_case_      (				self							)      ->				Union[str, Any]:
          pass







     @unittest.skip(reason="""Mask2Former does not have a get_input_embeddings method"""							)
     def 			snake_case_      (				self							)      ->				Tuple:
          pass







     @unittest.skip(reason="""Mask2Former is not a generative model"""							)
     def 			snake_case_      (				self							)      ->				List[Any]:
          pass







     @unittest.skip(reason="""Mask2Former does not use token embeddings"""							)
     def 			snake_case_      (				self							)      ->				Any:
          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							)      ->				Optional[int]:
          pass







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







     def 			snake_case_      (				self							)      ->				int:
          snake_case ,		snake_case			:      Optional[int]        =			self.model_tester.prepare_config_and_inputs_for_common()

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

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







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







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

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







     def 			snake_case_      (				self							)      ->				int:
          snake_case ,		snake_case			:      int        =			self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.create_and_check_maskaformer_model(__a							,**__a							,output_hidden_states=__a							)







     def 			snake_case_      (				self							)      ->				str:
          snake_case ,		snake_case			:      int        =			self.model_tester.prepare_config_and_inputs_for_common()

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







     def 			snake_case_      (				self							)      ->				Optional[int]:
          if not self.model_tester.is_training:
               return

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

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

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







     def 			snake_case_      (				self							)      ->				Optional[int]:
          snake_case			:      int        =			self.all_model_classes[1]
          snake_case ,		snake_case ,		snake_case ,		snake_case ,		snake_case			:      Tuple        =			self.model_tester.prepare_config_and_inputs()
          snake_case			:      Tuple        =			True
          snake_case			:      int        =			True

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

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

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

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

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

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

          outputs.loss.backward(retain_graph=__a							)

          self.assertIsNotNone(encoder_hidden_states.grad							)
          self.assertIsNotNone(pixel_decoder_hidden_states.grad							)
          self.assertIsNotNone(transformer_decoder_hidden_states.grad							)
          self.assertIsNotNone(attentions.grad							)
lowercase      :   Any      =       1E-4



def        lowerCamelCase__				(      ):
     snake_case			:      int        =			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""				)
     return image







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





     '''simple docstring'''



     @cached_property
     def 			snake_case_      (				self							)      ->				List[str]:
          return "facebook/mask2former-swin-small-coco-instance"







     @cached_property
     def 			snake_case_      (				self							)      ->				Optional[int]:
          return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints							) if is_vision_available() else None







     def 			snake_case_      (				self							)      ->				int:
          snake_case			:      int        =			MaskaFormerModel.from_pretrained(self.model_checkpoints							).to(__a							)
          snake_case			:      List[Any]        =			self.default_image_processor
          snake_case			:      Optional[Any]        =			prepare_img()
          snake_case			:      Optional[Any]        =			image_processor(__a							,return_tensors="""pt"""							).to(__a							)
          snake_case			:      int        =			inputs["""pixel_values"""].shape
          # check size is divisible by 32
          self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0							)
          # check size
          self.assertEqual(__a							,(1, 3, 384, 384)							)

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

          snake_case			:      List[Any]        =			torch.tensor(
              [[-0.2_790, -1.0_717, -1.1_668], [-0.5_128, -0.3_128, -0.4_987], [-0.5_832, 0.1_971, -0.0_197]]							).to(__a							)
          self.assertTrue(
              torch.allclose(
                  outputs.encoder_last_hidden_state[0, 0, :3, :3]							,__a							,atol=__a							)							)

          snake_case			:      List[str]        =			torch.tensor(
              [[0.8_973, 1.1_847, 1.1_776], [1.1_934, 1.5_040, 1.5_128], [1.1_153, 1.4_486, 1.4_951]]							).to(__a							)
          self.assertTrue(
              torch.allclose(
                  outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]							,__a							,atol=__a							)							)

          snake_case			:      Optional[Any]        =			torch.tensor(
              [[2.1_152, 1.7_000, -0.8_603], [1.5_808, 1.8_004, -0.9_353], [1.6_043, 1.7_495, -0.5_999]]							).to(__a							)
          self.assertTrue(
              torch.allclose(
                  outputs.transformer_decoder_last_hidden_state[0, :3, :3]							,__a							,atol=__a							)							)







     def 			snake_case_      (				self							)      ->				List[Any]:
          snake_case			:      List[Any]        =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints							).to(__a							).eval()
          snake_case			:      Optional[Any]        =			self.default_image_processor
          snake_case			:      List[Any]        =			prepare_img()
          snake_case			:      int        =			image_processor(__a							,return_tensors="""pt"""							).to(__a							)
          snake_case			:      Optional[Any]        =			inputs["""pixel_values"""].shape
          # check size is divisible by 32
          self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0							)
          # check size
          self.assertEqual(__a							,(1, 3, 384, 384)							)

          with torch.no_grad():
               snake_case			:      List[Any]        =			model(**__a							)
          # masks_queries_logits
          snake_case			:      str        =			outputs.masks_queries_logits
          self.assertEqual(
              masks_queries_logits.shape							,(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)							)
          snake_case			:      str        =			[
              [-8.7_839, -9.0_056, -8.8_121],
              [-7.4_104, -7.0_313, -6.5_401],
              [-6.6_105, -6.3_427, -6.4_675],
          ]
          snake_case			:      str        =			torch.tensor(__a							).to(__a							)
          self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]							,__a							,atol=__a							)							)
          # class_queries_logits
          snake_case			:      str        =			outputs.class_queries_logits
          self.assertEqual(class_queries_logits.shape							,(1, model.config.num_queries, model.config.num_labels + 1)							)
          snake_case			:      Optional[Any]        =			torch.tensor(
              [
                  [1.8_324, -8.0_835, -4.1_922],
                  [0.8_450, -9.0_050, -3.6_053],
                  [0.3_045, -7.7_293, -3.0_275],
              ]							).to(__a							)
          self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]							,__a							,atol=__a							)							)







     def 			snake_case_      (				self							)      ->				List[str]:
          snake_case			:      Union[str, Any]        =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints							).to(__a							).eval()
          snake_case			:      List[Any]        =			self.default_image_processor

          snake_case			:      Dict        =			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"""							,)

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

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

          self.assertTrue(outputs.loss is not None							)









'''simple docstring'''







# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings("""ignore""", category=UserWarning, module="""torch.optim.lr_scheduler""")







class 					_a :





     '''simple docstring'''



     def __init__(				self							,__a							,__a							,__a = True							,__a = False							)      ->				Dict:
          snake_case			:      List[Any]        =			scheduler
          snake_case			:      List[Any]        =			optimizers if isinstance(__a							,(list, tuple)							) else [optimizers]
          snake_case			:      List[str]        =			split_batches
          snake_case			:      List[str]        =			step_with_optimizer
          snake_case			:      int        =			GradientState()







     def 			snake_case_      (				self							,*__a							,**__a							)      ->				Optional[Any]:
          if not self.step_with_optimizer:
               # No link between scheduler and optimizer -> just step
               self.scheduler.step(*__a							,**__a							)
               return

          # Otherwise, first make sure the optimizer was stepped.
          if not self.gradient_state.sync_gradients:
               if self.gradient_state.adjust_scheduler:
                    self.scheduler._step_count += 1
               return

          for opt in self.optimizers:
               if opt.step_was_skipped:
                    return
          if self.split_batches:
               # Split batches -> the training dataloader batch size is not changed so one step per training step
               self.scheduler.step(*__a							,**__a							)
          else:
               # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
               # num_processes steps per training step
               snake_case			:      List[str]        =			AcceleratorState().num_processes
               for _ in range(__a							):
                    # Special case when using OneCycle and `drop_last` was not used
                    if hasattr(self.scheduler							,"""total_steps"""							):
                         if self.scheduler._step_count <= self.scheduler.total_steps:
                              self.scheduler.step(*__a							,**__a							)
                    else:
                         self.scheduler.step(*__a							,**__a							)







     def 			snake_case_      (				self							)      ->				Dict:
          return self.scheduler.get_last_lr()







     def 			snake_case_      (				self							)      ->				int:
          return self.scheduler.state_dict()







     def 			snake_case_      (				self							,__a							)      ->				Any:
          self.scheduler.load_state_dict(__a							)







     def 			snake_case_      (				self							)      ->				Any:
          return self.scheduler.get_lr()







     def 			snake_case_      (				self							,*__a							,**__a							)      ->				Optional[Any]:
          return self.scheduler.print_lr(*__a							,**__a							)



