



"""simple docstring"""







from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import BaseOutput, is_torch_available, is_transformers_available







@dataclass
class       UpperCamelCase__   (	lowercase__  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     42
      _SCREAMING_SNAKE_CASE  				=     42


if is_transformers_available() and is_torch_available():
     from .pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline







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





  """simple docstring"""

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


     '''simple docstring'''






     _a				:	Dict			=   parent
     _a				:	Union[str, Any]			=   batch_size
     _a				:	Tuple			=   is_training
     _a				:	List[str]			=   use_auxiliary_loss
     _a				:	Optional[Any]			=   num_queries
     _a				:	str			=   num_channels
     _a				:	List[str]			=   min_size
     _a				:	int			=   max_size
     _a				:	Optional[int]			=   num_labels
     _a				:	List[str]			=   hidden_dim
     _a				:	int			=   hidden_dim


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


     '''simple docstring'''






     _a				:	Tuple			=   floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]       ).to(
         _a       )

     _a				:	Optional[Any]			=   torch.ones([self.batch_size, self.min_size, self.max_size]							,device=_a       )

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

     _a				:	Dict			=   self.get_config()
     return config, pixel_values, pixel_mask, mask_labels, class_labels


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


     '''simple docstring'''






     _a				:	int			=   MaskaFormerConfig(
         hidden_size=self.hidden_dim							,)
     _a				:	str			=   self.num_queries
     _a				:	Union[str, Any]			=   self.num_labels

     _a				:	Tuple			=   [1, 1, 1, 1]
     _a				:	Dict			=   self.num_channels

     _a				:	str			=   64
     _a				:	Tuple			=   128
     _a				:	Optional[Any]			=   self.hidden_dim
     _a				:	Union[str, Any]			=   self.hidden_dim
     _a				:	List[Any]			=   self.hidden_dim
     return config


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     _a,  _a,  _a,  _a,  _a				:	Optional[Any]			=   self.prepare_config_and_inputs()
     _a				:	str			=   {'pixel_values': pixel_values, 'pixel_mask': pixel_mask}
     return config, inputs_dict


  def   __lowercase		(		self      :     List[str]							,_a      :     Optional[Any]							,_a      :     str       ):


     '''simple docstring'''






     _a				:	str			=   output.encoder_hidden_states
     _a				:	Any			=   output.pixel_decoder_hidden_states
     _a				:	Optional[Any]			=   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   __lowercase		(		self      :     List[str]							,_a      :     str							,_a      :     List[Any]							,_a      :     Any							,_a      :     Union[str, Any]=False       ):


     '''simple docstring'''






     with torch.no_grad():
        _a				:	str			=   MaskaFormerModel(config=_a       )
        model.to(_a       )
        model.eval()

        _a				:	Any			=   model(pixel_values=_a							,pixel_mask=_a       )
        _a				:	Optional[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   __lowercase		(		self      :     Tuple							,_a      :     List[Any]							,_a      :     Union[str, Any]							,_a      :     Tuple							,_a      :     List[str]							,_a      :     Any       ):


     '''simple docstring'''






     _a				:	int			=   MaskaFormerForUniversalSegmentation(config=_a       )
     model.to(_a       )
     model.eval()

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

        comm_check_on_output(_a       )

        _a				:	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 UpperCAmelCase__	( lowercase__    ,				lowercase__    ,				unittest.TestCase	):





  """simple docstring"""

  __UpperCAmelCase				:				Optional[int]   =					(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
  __UpperCAmelCase				:				Dict   =					{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

  __UpperCAmelCase				:				Dict   =					False
  __UpperCAmelCase				:				Tuple   =					False
  __UpperCAmelCase				:				Dict   =					False
  __UpperCAmelCase				:				List[Any]   =					False


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   MaskaFormerModelTester(self       )
     _a				:	Dict			=   ConfigTester(self							,config_class=_a							,has_text_modality=_a       )


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     self.config_tester.run_common_tests()


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     _a,  _a				:	List[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   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     _a				:	Dict			=   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   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     pass


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


     '''simple docstring'''






     pass


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


     '''simple docstring'''






     pass


  @unittest.skip(reason='Mask2Former does not use token embeddings'       )
  def   __lowercase		(		self      :     Optional[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   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     pass


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


     '''simple docstring'''






     pass


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a,  _a				:	Any			=   self.model_tester.prepare_config_and_inputs_for_common()

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

        _a				:	List[Any]			=   ['pixel_values']
        self.assertListEqual(arg_names[:1]							,_a       )


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


     '''simple docstring'''






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


  def   __lowercase		(		self      :     List[Any]       ):


     '''simple docstring'''






     _a				:	int			=   (self.model_tester.min_size,) * 2
     _a				:	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(),
     }
     _a				:	List[Any]			=   self.model_tester.get_config()

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


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a,  _a				:	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   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a,  _a				:	Dict			=   self.model_tester.prepare_config_and_inputs_for_common()

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


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     if not self.model_tester.is_training:
        return

     _a				:	List[str]			=   self.all_model_classes[1]
     _a,  _a,  _a,  _a,  _a				:	List[str]			=   self.model_tester.prepare_config_and_inputs()

     _a				:	Any			=   model_class(_a       )
     model.to(_a       )
     model.train()

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


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a				:	int			=   self.all_model_classes[1]
     _a,  _a,  _a,  _a,  _a				:	List[Any]			=   self.model_tester.prepare_config_and_inputs()
     _a				:	str			=   True
     _a				:	str			=   True

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

     _a				:	Optional[int]			=   model(_a							,mask_labels=_a							,class_labels=_a       )

     _a				:	Tuple			=   outputs.encoder_hidden_states[0]
     encoder_hidden_states.retain_grad()

     _a				:	str			=   outputs.pixel_decoder_hidden_states[0]
     pixel_decoder_hidden_states.retain_grad()

     _a				:	Dict			=   outputs.transformer_decoder_hidden_states[0]
     transformer_decoder_hidden_states.retain_grad()

     _a				:	List[str]			=   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         =    1e-4





def 					UpperCAmelCase_		():






   """simple docstring"""

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







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





  """simple docstring"""

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


     '''simple docstring'''






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


  @cached_property
  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






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


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	List[str]			=   MaskaFormerModel.from_pretrained(self.model_checkpoints       ).to(_a       )
     _a				:	int			=   self.default_image_processor
     _a				:	Tuple			=   prepare_img()
     _a				:	Any			=   image_processor(_a							,return_tensors='pt'       ).to(_a       )
     _a				:	Union[str, 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():
        _a				:	Optional[Any]			=   model(**_a       )

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

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

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


  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     _a				:	List[Any]			=   MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints       ).to(_a       ).eval()
     _a				:	Optional[Any]			=   self.default_image_processor
     _a				:	List[Any]			=   prepare_img()
     _a				:	str			=   image_processor(_a							,return_tensors='pt'       ).to(_a       )
     _a				:	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():
        _a				:	Optional[int]			=   model(**_a       )
     # masks_queries_logits
     _a				:	Dict			=   outputs.masks_queries_logits
     self.assertEqual(
         masks_queries_logits.shape							,(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)       )
     _a				:	Dict			=   [
         [-8.7839, -9.0056, -8.8121],
         [-7.4104, -7.0313, -6.5401],
         [-6.6105, -6.3427, -6.4675],
     ]
     _a				:	Optional[Any]			=   torch.tensor(_a       ).to(_a       )
     self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]							,_a							,atol=_a       )       )
     # class_queries_logits
     _a				:	str			=   outputs.class_queries_logits
     self.assertEqual(class_queries_logits.shape							,(1, model.config.num_queries, model.config.num_labels + 1)       )
     _a				:	str			=   torch.tensor(
         [
             [1.8324, -8.0835, -4.1922],
             [0.8450, -9.0050, -3.6053],
             [0.3045, -7.7293, -3.0275],
         ]       ).to(_a       )
     self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]							,_a							,atol=_a       )       )


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     _a				:	Any			=   MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints       ).to(_a       ).eval()
     _a				:	Tuple			=   self.default_image_processor

     _a				:	Tuple			=   image_processor(
         [np.zeros((3, 800, 1333)       ), np.zeros((3, 800, 1333)       )]							,segmentation_maps=[np.zeros((384, 384)       ).astype(np.floataa       ), np.zeros((384, 384)       ).astype(np.floataa       )]							,return_tensors='pt'							,)

     _a				:	str			=   inputs['pixel_values'].to(_a       )
     _a				:	str			=   [el.to(_a       ) for el in inputs['mask_labels']]
     _a				:	Dict			=   [el.to(_a       ) for el in inputs['class_labels']]

     with torch.no_grad():
        _a				:	List[str]			=   model(**_a       )

     self.assertTrue(outputs.loss is not None       )



