





"""simple docstring"""






import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import DetrConfig, MaskFormerConfig, SwinConfig, 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 MaskFormerForInstanceSegmentation, MaskFormerModel

    if is_vision_available():
        from transformers import MaskFormerImageProcessor

if is_vision_available():
    from PIL import Image



class 	lowerCAmelCase					:

       '''simple docstring'''





       def __init__(   self    ,  a__    ,  a__=2    ,  a__=True    ,  a__=False    ,  a__=10    ,  a__=3    ,  a__=32 * 4    ,  a__=32 * 6    ,  a__=4    ,  a__=32    ,  ):
              _UpperCAmelCase     =     parent
              _UpperCAmelCase     =     batch_size
              _UpperCAmelCase     =     is_training
              _UpperCAmelCase     =     use_auxiliary_loss
              _UpperCAmelCase     =     num_queries
              _UpperCAmelCase     =     num_channels
              _UpperCAmelCase     =     min_size
              _UpperCAmelCase     =     max_size
              _UpperCAmelCase     =     num_labels
              _UpperCAmelCase     =     mask_feature_size






       def 							__A (   self     ):
              _UpperCAmelCase     =     floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]     ).to(
                  __lowerCAmelCase     )

              _UpperCAmelCase     =     torch.ones([self.batch_size, self.min_size, self.max_size]    ,  device=__lowerCAmelCase     )

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

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






       def 							__A (   self     ):
              return MaskFormerConfig.from_backbone_and_decoder_configs(
                  backbone_config=SwinConfig(
                      depths=[1, 1, 1, 1]    ,  )    ,  decoder_config=DetrConfig(
                      decoder_ffn_dim=1_28    ,  num_queries=self.num_queries    ,  decoder_attention_heads=2    ,  d_model=self.mask_feature_size    ,  )    ,  mask_feature_size=self.mask_feature_size    ,  fpn_feature_size=self.mask_feature_size    ,  num_channels=self.num_channels    ,  num_labels=self.num_labels    ,  )






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






       def 							__A (   self    ,  a__    ,  a__     ):
              _UpperCAmelCase     =     output.encoder_hidden_states
              _UpperCAmelCase     =     output.pixel_decoder_hidden_states
              _UpperCAmelCase     =     output.transformer_decoder_hidden_states

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






       def 							__A (   self    ,  a__    ,  a__    ,  a__    ,  a__=False     ):






              with torch.no_grad():
                     _UpperCAmelCase     =     MaskFormerModel(config=__lowerCAmelCase     )
                     model.to(__lowerCAmelCase     )
                     model.eval()

                     _UpperCAmelCase     =     model(pixel_values=__lowerCAmelCase    ,  pixel_mask=__lowerCAmelCase     )
                     _UpperCAmelCase     =     model(__lowerCAmelCase    ,  output_hidden_states=__lowerCAmelCase     )
              # the correct shape of output.transformer_decoder_hidden_states ensure the correcteness of the
              # encoder and pixel decoder
              self.parent.assertEqual(
                  output.transformer_decoder_last_hidden_state.shape    ,  (self.batch_size, self.num_queries, self.mask_feature_size)    ,  )
              # 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(__lowerCAmelCase    ,  __lowerCAmelCase     )






       def 							__A (   self    ,  a__    ,  a__    ,  a__    ,  a__    ,  a__     ):
              _UpperCAmelCase     =     MaskFormerForInstanceSegmentation(config=__lowerCAmelCase     )
              model.to(__lowerCAmelCase     )
              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():
                     _UpperCAmelCase     =     model(pixel_values=__lowerCAmelCase    ,  pixel_mask=__lowerCAmelCase     )
                     _UpperCAmelCase     =     model(__lowerCAmelCase     )

                     comm_check_on_output(__lowerCAmelCase     )

                     _UpperCAmelCase     =     model(
                         pixel_values=__lowerCAmelCase    ,  pixel_mask=__lowerCAmelCase    ,  mask_labels=__lowerCAmelCase    ,  class_labels=__lowerCAmelCase     )

              comm_check_on_output(__lowerCAmelCase     )

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



@require_torch
class 	lowerCAmelCase					(      UpperCAmelCase__			,	UpperCAmelCase__			,	unittest.TestCase							):

       '''simple docstring'''





       lowerCAmelCase__						 =  (MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
       lowerCAmelCase__						 =  (
           {"""feature-extraction""": MaskFormerModel, """image-segmentation""": MaskFormerForInstanceSegmentation}
           if is_torch_available()
           else {}
       )

       lowerCAmelCase__						 =  False
       lowerCAmelCase__						 =  False
       lowerCAmelCase__						 =  False
       lowerCAmelCase__						 =  False






       def 							__A (   self     ):
              _UpperCAmelCase     =     MaskFormerModelTester(self     )
              _UpperCAmelCase     =     ConfigTester(self    ,  config_class=__lowerCAmelCase    ,  has_text_modality=__lowerCAmelCase     )






       def 							__A (   self     ):
              self.config_tester.run_common_tests()






       def 							__A (   self     ):
              _UpperCAmelCase    ,							_UpperCAmelCase     =     self.model_tester.prepare_config_and_inputs_for_common()
              self.model_tester.create_and_check_maskformer_model(__lowerCAmelCase    ,  **__lowerCAmelCase    ,  output_hidden_states=__lowerCAmelCase     )






       def 							__A (   self     ):
              _UpperCAmelCase     =     self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*__lowerCAmelCase     )






       @unittest.skip(reason='MaskFormer does not use inputs_embeds'     )
       def 							__A (   self     ):
              pass






       @unittest.skip(reason='MaskFormer does not have a get_input_embeddings method'     )
       def 							__A (   self     ):
              pass






       @unittest.skip(reason='MaskFormer is not a generative model'     )
       def 							__A (   self     ):
              pass






       @unittest.skip(reason='MaskFormer does not use token embeddings'     )
       def 							__A (   self     ):
              pass






       @require_torch_multi_gpu
       @unittest.skip(
           reason='MaskFormer has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'     )
       def 							__A (   self     ):
              pass






       @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'     )
       def 							__A (   self     ):
              pass






       def 							__A (   self     ):
              _UpperCAmelCase    ,							_UpperCAmelCase     =     self.model_tester.prepare_config_and_inputs_for_common()

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

                     _UpperCAmelCase     =     ['pixel_values']
                     self.assertListEqual(arg_names[:1]    ,  __lowerCAmelCase     )






       @slow
       def 							__A (   self     ):






              for model_name in ["facebook/maskformer-swin-small-coco"]:
                     _UpperCAmelCase     =     MaskFormerModel.from_pretrained(__lowerCAmelCase     )
                     self.assertIsNotNone(__lowerCAmelCase     )






       def 							__A (   self     ):
              _UpperCAmelCase     =     (self.model_tester.min_size,) * 2
              _UpperCAmelCase     =     {
                  'pixel_values': torch.randn((2, 3, *size)    ,  device=__lowerCAmelCase     ),
                  'mask_labels': torch.randn((2, 10, *size)    ,  device=__lowerCAmelCase     ),
                  'class_labels': torch.zeros(2    ,  10    ,  device=__lowerCAmelCase     ).long(),
              }

              _UpperCAmelCase     =     MaskFormerForInstanceSegmentation(MaskFormerConfig()     ).to(__lowerCAmelCase     )
              _UpperCAmelCase     =     model(**__lowerCAmelCase     )
              self.assertTrue(outputs.loss is not None     )






       def 							__A (   self     ):
              _UpperCAmelCase    ,							_UpperCAmelCase     =     self.model_tester.prepare_config_and_inputs_for_common()
              self.model_tester.create_and_check_maskformer_model(__lowerCAmelCase    ,  **__lowerCAmelCase    ,  output_hidden_states=__lowerCAmelCase     )






       def 							__A (   self     ):
              _UpperCAmelCase    ,							_UpperCAmelCase     =     self.model_tester.prepare_config_and_inputs_for_common()

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






       def 							__A (   self     ):






              if not self.model_tester.is_training:
                     return
              # only MaskFormerForInstanceSegmentation has the loss
              _UpperCAmelCase     =     self.all_model_classes[1]
              _UpperCAmelCase    ,							_UpperCAmelCase    ,							_UpperCAmelCase    ,							_UpperCAmelCase    ,							_UpperCAmelCase     =     self.model_tester.prepare_config_and_inputs()

              _UpperCAmelCase     =     model_class(__lowerCAmelCase     )
              model.to(__lowerCAmelCase     )
              model.train()

              _UpperCAmelCase     =     model(__lowerCAmelCase    ,  mask_labels=__lowerCAmelCase    ,  class_labels=__lowerCAmelCase     ).loss
              loss.backward()






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

              _UpperCAmelCase     =     model_class(__lowerCAmelCase     )
              model.to(__lowerCAmelCase     )
              model.train()

              _UpperCAmelCase     =     model(__lowerCAmelCase    ,  mask_labels=__lowerCAmelCase    ,  class_labels=__lowerCAmelCase     )

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

              _UpperCAmelCase     =     outputs.pixel_decoder_hidden_states[0]
              pixel_decoder_hidden_states.retain_grad()
              # we requires_grad=True in inputs_embeds (line 2152), the original implementation don't
              _UpperCAmelCase     =     outputs.transformer_decoder_hidden_states[0]
              transformer_decoder_hidden_states.retain_grad()

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

              outputs.loss.backward(retain_graph=__lowerCAmelCase     )

              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       __lowerCamelCase    (       )		->       Dict:






       """simple docstring"""




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



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

       '''simple docstring'''





       @cached_property
       def 							__A (   self     ):
              return (
                  MaskFormerImageProcessor.from_pretrained('facebook/maskformer-swin-small-coco'     )
                  if is_vision_available()
                  else None
              )






       def 							__A (   self     ):
              _UpperCAmelCase     =     MaskFormerModel.from_pretrained('facebook/maskformer-swin-small-coco'     ).to(__lowerCAmelCase     )
              _UpperCAmelCase     =     self.default_image_processor
              _UpperCAmelCase     =     prepare_img()
              _UpperCAmelCase     =     image_processor(__lowerCAmelCase    ,  return_tensors='pt'     ).to(__lowerCAmelCase     )
              _UpperCAmelCase     =     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(__lowerCAmelCase    ,  (1, 3, 8_00, 10_88)     )

              with torch.no_grad():
                     _UpperCAmelCase     =     model(**__lowerCAmelCase     )

              _UpperCAmelCase     =     torch.tensor(
                  [[-0.0_482, 0.9_228, 0.4_951], [-0.2_547, 0.8_017, 0.8_527], [-0.0_069, 0.3_385, -0.0_089]]     ).to(__lowerCAmelCase     )
              self.assertTrue(
                  torch.allclose(
                      outputs.encoder_last_hidden_state[0, 0, :3, :3]    ,  __lowerCAmelCase    ,  atol=__lowerCAmelCase     )     )

              _UpperCAmelCase     =     torch.tensor(
                  [[-0.8_422, -0.8_434, -0.9_718], [-1.0_144, -0.5_565, -0.4_195], [-1.0_038, -0.4_484, -0.1_961]]     ).to(__lowerCAmelCase     )
              self.assertTrue(
                  torch.allclose(
                      outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]    ,  __lowerCAmelCase    ,  atol=__lowerCAmelCase     )     )

              _UpperCAmelCase     =     torch.tensor(
                  [[0.2_852, -0.0_159, 0.9_735], [0.6_254, 0.1_858, 0.8_529], [-0.0_680, -0.4_116, 1.8_413]]     ).to(__lowerCAmelCase     )
              self.assertTrue(
                  torch.allclose(
                      outputs.transformer_decoder_last_hidden_state[0, :3, :3]    ,  __lowerCAmelCase    ,  atol=__lowerCAmelCase     )     )






       def 							__A (   self     ):
              _UpperCAmelCase     =     (
                  MaskFormerForInstanceSegmentation.from_pretrained('facebook/maskformer-swin-small-coco'     )
                  .to(__lowerCAmelCase     )
                  .eval()
              )
              _UpperCAmelCase     =     self.default_image_processor
              _UpperCAmelCase     =     prepare_img()
              _UpperCAmelCase     =     image_processor(__lowerCAmelCase    ,  return_tensors='pt'     ).to(__lowerCAmelCase     )
              _UpperCAmelCase     =     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(__lowerCAmelCase    ,  (1, 3, 8_00, 10_88)     )

              with torch.no_grad():
                     _UpperCAmelCase     =     model(**__lowerCAmelCase     )
              # masks_queries_logits
              _UpperCAmelCase     =     outputs.masks_queries_logits
              self.assertEqual(
                  masks_queries_logits.shape    ,  (1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)    ,  )
              _UpperCAmelCase     =     [
                  [-1.3_737_124, -1.7_724_937, -1.9_364_233],
                  [-1.5_977_281, -1.9_867_939, -2.1_523_695],
                  [-1.5_795_398, -1.9_269_832, -2.093_942],
              ]
              _UpperCAmelCase     =     torch.tensor(__lowerCAmelCase     ).to(__lowerCAmelCase     )
              self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]    ,  __lowerCAmelCase    ,  atol=__lowerCAmelCase     )     )
              # class_queries_logits
              _UpperCAmelCase     =     outputs.class_queries_logits
              self.assertEqual(
                  class_queries_logits.shape    ,  (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)     )
              _UpperCAmelCase     =     torch.tensor(
                  [
                      [1.6512E00, -5.2572E00, -3.3519E00],
                      [3.6169E-02, -5.9025E00, -2.9313E00],
                      [1.0766E-04, -7.7630E00, -5.1263E00],
                  ]     ).to(__lowerCAmelCase     )
              self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]    ,  __lowerCAmelCase    ,  atol=__lowerCAmelCase     )     )






       def 							__A (   self     ):
              _UpperCAmelCase     =     (
                  MaskFormerForInstanceSegmentation.from_pretrained('facebook/maskformer-resnet101-coco-stuff'     )
                  .to(__lowerCAmelCase     )
                  .eval()
              )
              _UpperCAmelCase     =     self.default_image_processor
              _UpperCAmelCase     =     prepare_img()
              _UpperCAmelCase     =     image_processor(__lowerCAmelCase    ,  return_tensors='pt'     ).to(__lowerCAmelCase     )
              _UpperCAmelCase     =     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(__lowerCAmelCase    ,  (1, 3, 8_00, 10_88)     )

              with torch.no_grad():
                     _UpperCAmelCase     =     model(**__lowerCAmelCase     )
              # masks_queries_logits
              _UpperCAmelCase     =     outputs.masks_queries_logits
              self.assertEqual(
                  masks_queries_logits.shape    ,  (1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)    ,  )
              _UpperCAmelCase     =     [[-0.9_046, -2.6_366, -4.6_062], [-3.4_179, -5.7_890, -8.8_057], [-4.9_179, -7.6_560, -10.77_11]]
              _UpperCAmelCase     =     torch.tensor(__lowerCAmelCase     ).to(__lowerCAmelCase     )
              self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]    ,  __lowerCAmelCase    ,  atol=__lowerCAmelCase     )     )
              # class_queries_logits
              _UpperCAmelCase     =     outputs.class_queries_logits
              self.assertEqual(
                  class_queries_logits.shape    ,  (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)     )
              _UpperCAmelCase     =     torch.tensor(
                  [[4.7_188, -3.2_585, -2.8_857], [6.6_871, -2.9_181, -1.2_487], [7.2_449, -2.2_764, -2.1_874]]     ).to(__lowerCAmelCase     )
              self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]    ,  __lowerCAmelCase    ,  atol=__lowerCAmelCase     )     )






       def 							__A (   self     ):
              _UpperCAmelCase     =     (
                  MaskFormerForInstanceSegmentation.from_pretrained('facebook/maskformer-swin-small-coco'     )
                  .to(__lowerCAmelCase     )
                  .eval()
              )
              _UpperCAmelCase     =     self.default_image_processor

              _UpperCAmelCase     =     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'    ,  )

              _UpperCAmelCase     =     inputs['pixel_values'].to(__lowerCAmelCase     )
              _UpperCAmelCase     =     [el.to(__lowerCAmelCase     ) for el in inputs['mask_labels']]
              _UpperCAmelCase     =     [el.to(__lowerCAmelCase     ) for el in inputs['class_labels']]

              with torch.no_grad():
                     _UpperCAmelCase     =     model(**__lowerCAmelCase     )

              self.assertTrue(outputs.loss is not None     )













"""simple docstring"""






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

import numpy as np
import PIL

from ...utils import BaseOutput, OptionalDependencyNotAvailable, is_torch_available, is_transformers_available
from .timesteps import (
    fastaa_timesteps,
    smartaa_timesteps,
    smartaa_timesteps,
    smartaaa_timesteps,
    smartaaa_timesteps,
    superaa_timesteps,
    superaa_timesteps,
    superaaa_timesteps,
)



@dataclass
class 	lowerCAmelCase					(      snake_case							):
       lowerCAmelCase__						 =  42
       lowerCAmelCase__						 =  42
       lowerCAmelCase__						 =  42


try:
    if not (is_transformers_available() and is_torch_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
    from .pipeline_if import IFPipeline
    from .pipeline_if_imgaimg import IFImgaImgPipeline
    from .pipeline_if_imgaimg_superresolution import IFImgaImgSuperResolutionPipeline
    from .pipeline_if_inpainting import IFInpaintingPipeline
    from .pipeline_if_inpainting_superresolution import IFInpaintingSuperResolutionPipeline
    from .pipeline_if_superresolution import IFSuperResolutionPipeline
    from .safety_checker import IFSafetyChecker
    from .watermark import IFWatermarker






