



"""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	:      Union[str, Any]	,      A	:      Union[str, Any]	,      A	:      int=2	,      A	:      Optional[Any]=True	,      A	:      List[Any]=False	,      A	:      Union[str, Any]=1_0	,      A	:      Optional[int]=3	,      A	:      Optional[Any]=3_2 * 8	,      A	:      Dict=3_2 * 8	,      A	:      str=4	,      A	:      Tuple=6_4	,      ):
           _UpperCAmelCase     :      str														=       parent
           _UpperCAmelCase     :      List[Any]														=       batch_size
           _UpperCAmelCase     :      int														=       is_training
           _UpperCAmelCase     :      Optional[Any]														=       use_auxiliary_loss
           _UpperCAmelCase     :      Tuple														=       num_queries
           _UpperCAmelCase     :      Union[str, Any]														=       num_channels
           _UpperCAmelCase     :      Any														=       min_size
           _UpperCAmelCase     :      Dict														=       max_size
           _UpperCAmelCase     :      Optional[Any]														=       num_labels
           _UpperCAmelCase     :      List[str]														=       hidden_dim
           _UpperCAmelCase     :      List[Any]														=       hidden_dim



      def 		snake_case_   (	self	:      Optional[int]		):
           _UpperCAmelCase     :      List[str]														=       floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]		).to(
               A		)

           _UpperCAmelCase     :      str														=       torch.ones([self.batch_size, self.min_size, self.max_size]	,      device=A		)

           _UpperCAmelCase     :      Any														=       (
               torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]	,      device=A		) > 0.5
           ).float()
           _UpperCAmelCase     :      Dict														=       (torch.rand((self.batch_size, self.num_labels)	,      device=A		) > 0.5).long()

           _UpperCAmelCase     :      Tuple														=       self.get_config()
           return config, pixel_values, pixel_mask, mask_labels, class_labels



      def 		snake_case_   (	self	:      str		):
           _UpperCAmelCase     :      Optional[Any]														=       MaskaFormerConfig(
               hidden_size=self.hidden_dim	,      )
           _UpperCAmelCase     :      Dict														=       self.num_queries
           _UpperCAmelCase     :      List[str]														=       self.num_labels

           _UpperCAmelCase     :      List[str]														=       [1, 1, 1, 1]
           _UpperCAmelCase     :      Dict														=       self.num_channels

           _UpperCAmelCase     :      str														=       6_4
           _UpperCAmelCase     :      List[Any]														=       1_2_8
           _UpperCAmelCase     :      List[str]														=       self.hidden_dim
           _UpperCAmelCase     :      Optional[Any]														=       self.hidden_dim
           _UpperCAmelCase     :      Optional[Any]														=       self.hidden_dim
           return config



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



      def 		snake_case_   (	self	:      str	,      A	:      Optional[int]	,      A	:      Optional[Any]		):
           _UpperCAmelCase     :      str														=       output.encoder_hidden_states
           _UpperCAmelCase     :      Any														=       output.pixel_decoder_hidden_states
           _UpperCAmelCase     :      Optional[int]														=       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	:      int	,      A	:      Any	,      A	:      Optional[int]	,      A	:      Any	,      A	:      Optional[int]=False		):
           with torch.no_grad():
                _UpperCAmelCase     :      Any														=       MaskaFormerModel(config=A		)
                model.to(A		)
                model.eval()

                _UpperCAmelCase     :      Tuple														=       model(pixel_values=A	,      pixel_mask=A		)
                _UpperCAmelCase     :      Dict														=       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	:      List[str]	,      A	:      Dict	,      A	:      Optional[Any]	,      A	:      Optional[int]	,      A	:      Dict	,      A	:      Dict		):
           _UpperCAmelCase     :      Union[str, Any]														=       MaskaFormerForUniversalSegmentation(config=A		)
           model.to(A		)
           model.eval()

           def comm_check_on_output(A	:      Dict		):
                # 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     :      Dict														=       model(pixel_values=A	,      pixel_mask=A		)
                _UpperCAmelCase     :      List[str]														=       model(A		)

                comm_check_on_output(A		)

                _UpperCAmelCase     :      Union[str, Any]														=       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       UpperCAmelCase_    (							_UpperCamelCase			,     _UpperCamelCase			,     unittest.TestCase		):
      __SCREAMING_SNAKE_CASE		:  int						=					(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
      __SCREAMING_SNAKE_CASE		:  int						=					{'feature-extraction': MaskaFormerModel} if is_torch_available() else {}

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



      def 		snake_case_   (	self	:      Dict		):
           _UpperCAmelCase     :      Optional[int]														=       MaskaFormerModelTester(self		)
           _UpperCAmelCase     :      List[str]														=       ConfigTester(self	,      config_class=A	,      has_text_modality=A		)



      def 		snake_case_   (	self	:      Optional[Any]		):
           self.config_tester.run_common_tests()



      def 		snake_case_   (	self	:      Tuple		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      str														=       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	:      Optional[int]		):
           _UpperCAmelCase     :      Optional[int]														=       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	:      List[str]		):
           pass



      @unittest.skip(reason="Mask2Former does not have a get_input_embeddings method"		)
      def 		snake_case_   (	self	:      Optional[int]		):
           pass



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



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



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



      def 		snake_case_   (	self	:      List[Any]		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      str														=       self.model_tester.prepare_config_and_inputs_for_common()

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

                _UpperCAmelCase     :      Union[str, Any]														=       ["pixel_values"]
                self.assertListEqual(arg_names[:1]	,      A		)



      @slow
      def 		snake_case_   (	self	:      Tuple		):
           for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
                _UpperCAmelCase     :      Dict														=       MaskaFormerModel.from_pretrained(A		)
                self.assertIsNotNone(A		)



      def 		snake_case_   (	self	:      List[str]		):
           _UpperCAmelCase     :      List[Any]														=       (self.model_tester.min_size,) * 2
           _UpperCAmelCase     :      Optional[int]														=       {
               "pixel_values": torch.randn((2, 3, *size)	,      device=A		),
               "mask_labels": torch.randn((2, 1_0, *size)	,      device=A		),
               "class_labels": torch.zeros(2	,      1_0	,      device=A		).long(),
           }
           _UpperCAmelCase     :      List[Any]														=       self.model_tester.get_config()

           _UpperCAmelCase     :      int														=       MaskaFormerForUniversalSegmentation(A		).to(A		)
           _UpperCAmelCase     :      str														=       model(**A		)
           self.assertTrue(outputs.loss is not None		)



      def 		snake_case_   (	self	:      Dict		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      List[Any]														=       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	:      List[Any]		):
           _UpperCAmelCase							,			_UpperCAmelCase     :      Tuple														=       self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                _UpperCAmelCase     :      Dict														=       model_class(A		).to(A		)
                _UpperCAmelCase     :      Optional[int]														=       model(**A	,      output_attentions=A		)
                self.assertTrue(outputs.attentions is not None		)



      def 		snake_case_   (	self	:      Union[str, Any]		):
           if not self.model_tester.is_training:
                return

           _UpperCAmelCase     :      str														=       self.all_model_classes[1]
           _UpperCAmelCase							,			_UpperCAmelCase							,			_UpperCAmelCase							,			_UpperCAmelCase							,			_UpperCAmelCase     :      Dict														=       self.model_tester.prepare_config_and_inputs()

           _UpperCAmelCase     :      int														=       model_class(A		)
           model.to(A		)
           model.train()

           _UpperCAmelCase     :      Union[str, Any]														=       model(A	,      mask_labels=A	,      class_labels=A		).loss
           loss.backward()



      def 		snake_case_   (	self	:      List[Any]		):
           _UpperCAmelCase     :      Tuple														=       self.all_model_classes[1]
           _UpperCAmelCase							,			_UpperCAmelCase							,			_UpperCAmelCase							,			_UpperCAmelCase							,			_UpperCAmelCase     :      str														=       self.model_tester.prepare_config_and_inputs()
           _UpperCAmelCase     :      int														=       True
           _UpperCAmelCase     :      Optional[Any]														=       True

           _UpperCAmelCase     :      str														=       model_class(A		).to(A		)
           model.train()

           _UpperCAmelCase     :      Optional[Any]														=       model(A	,      mask_labels=A	,      class_labels=A		)

           _UpperCAmelCase     :      List[Any]														=       outputs.encoder_hidden_states[0]
           encoder_hidden_states.retain_grad()

           _UpperCAmelCase     :      Tuple														=       outputs.pixel_decoder_hidden_states[0]
           pixel_decoder_hidden_states.retain_grad()

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

           _UpperCAmelCase     :      Tuple														=       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		)







_lowerCAmelCase							:   Optional[int]				=				1E-4

def 		__snake_case		( )    ->  Any:





     '''simple docstring'''

     _UpperCAmelCase     :      Optional[int]														=       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"							)
     return image






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



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



      @cached_property
      def 		snake_case_   (	self	:      Any		):
           return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints		) if is_vision_available() else None



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

           with torch.no_grad():
                _UpperCAmelCase     :      Optional[int]														=       model(**A		)

           _UpperCAmelCase     :      Dict														=       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		)		)

           _UpperCAmelCase     :      int														=       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		)		)

           _UpperCAmelCase     :      Optional[int]														=       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	:      Tuple		):
           _UpperCAmelCase     :      Union[str, Any]														=       MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints		).to(A		).eval()
           _UpperCAmelCase     :      Union[str, Any]														=       self.default_image_processor
           _UpperCAmelCase     :      Optional[int]														=       prepare_img()
           _UpperCAmelCase     :      Union[str, Any]														=       image_processor(A	,      return_tensors="pt"		).to(A		)
           _UpperCAmelCase     :      Union[str, Any]														=       inputs["pixel_values"].shape
           # check size is divisible by 32
           self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0		)
           # check size
           self.assertEqual(A	,      (1, 3, 3_8_4, 3_8_4)		)

           with torch.no_grad():
                _UpperCAmelCase     :      str														=       model(**A		)
           # masks_queries_logits
           _UpperCAmelCase     :      int														=       outputs.masks_queries_logits
           self.assertEqual(
               masks_queries_logits.shape	,      (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)		)
           _UpperCAmelCase     :      int														=       [
               [-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],
           ]
           _UpperCAmelCase     :      Any														=       torch.tensor(A		).to(A		)
           self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]	,      A	,      atol=A		)		)
           # class_queries_logits
           _UpperCAmelCase     :      List[str]														=       outputs.class_queries_logits
           self.assertEqual(class_queries_logits.shape	,      (1, model.config.num_queries, model.config.num_labels + 1)		)
           _UpperCAmelCase     :      List[str]														=       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	:      Optional[int]		):
           _UpperCAmelCase     :      List[Any]														=       MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints		).to(A		).eval()
           _UpperCAmelCase     :      Tuple														=       self.default_image_processor

           _UpperCAmelCase     :      List[str]														=       image_processor(
               [np.zeros((3, 8_0_0, 1_3_3_3)		), np.zeros((3, 8_0_0, 1_3_3_3)		)]	,      segmentation_maps=[np.zeros((3_8_4, 3_8_4)		).astype(np.floataa		), np.zeros((3_8_4, 3_8_4)		).astype(np.floataa		)]	,      return_tensors="pt"	,      )

           _UpperCAmelCase     :      Optional[Any]														=       inputs["pixel_values"].to(A		)
           _UpperCAmelCase     :      str														=       [el.to(A		) for el in inputs["mask_labels"]]
           _UpperCAmelCase     :      Optional[Any]														=       [el.to(A		) for el in inputs["class_labels"]]

           with torch.no_grad():
                _UpperCAmelCase     :      Tuple														=       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       UpperCAmelCase_    :



      def __init__(	self	:      Optional[Any]	,      A	:      List[str]	,      A	:      List[str]	,      A	:      bool = True	,      A	:      bool = False		):
           _UpperCAmelCase     :      Dict														=       scheduler
           _UpperCAmelCase     :      str														=       optimizers if isinstance(A	,      (list, tuple)		) else [optimizers]
           _UpperCAmelCase     :      List[str]														=       split_batches
           _UpperCAmelCase     :      Dict														=       step_with_optimizer
           _UpperCAmelCase     :      Tuple														=       GradientState()



      def 		snake_case_   (	self	:      Optional[int]	,      *A	:      List[Any]	,      **A	:      Optional[int]		):
           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
                _UpperCAmelCase     :      int														=       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	:      Union[str, Any]		):
           return self.scheduler.get_last_lr()



      def 		snake_case_   (	self	:      str		):
           return self.scheduler.state_dict()



      def 		snake_case_   (	self	:      str	,      A	:      Optional[Any]		):
           self.scheduler.load_state_dict(A		)



      def 		snake_case_   (	self	:      Tuple		):
           return self.scheduler.get_lr()



      def 		snake_case_   (	self	:      Optional[int]	,      *A	:      List[str]	,      **A	:      Optional[Any]		):
           return self.scheduler.print_lr(*A	,      **A		)
