codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch maskformer model import copy import unittest import numpy as np from tests testmodelingcommon import floatstensor from transformers import detrconfig maskformerconfig swinconfig istorchavailable isvisionavailable from transformers testingutils import requiretorch requiretorchaccelerator requiretorchfp16 requiretorchmultigpu requirevision slow torchdevice from transformers utils import cachedproperty from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from transformers import maskformerforinstancesegmentation maskformermodel if isvisionavailable from transformers import maskformerimageprocessor if isvisionavailable from pil import image class maskformermodeltester def init self parent batchsize2 istrainingtrue useauxiliarylossfalse numqueries10 numchannels3 minsize32 4 maxsize32 6 numlabels4 maskfeaturesize32 numhiddenlayers2 numattentionheads2 self parent parent self batchsize batchsize self istraining istraining self useauxiliaryloss useauxiliaryloss self numqueries numqueries self numchannels numchannels self minsize minsize self maxsize maxsize self numlabels numlabels self maskfeaturesize maskfeaturesize this is passed to the decoder config we add it to the model tester here for testing self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads def prepareconfigandinputsself pixelvalues floatstensorself batchsize self numchannels self minsize self maxsize to torchdevice pixelmask torch onesself batchsize self minsize self maxsize devicetorchdevice masklabels torch randself batchsize self numlabels self minsize self maxsize devicetorchdevice 0 5 float classlabels torch randself batchsize self numlabels devicetorchdevice 0 5 long config self getconfig return config pixelvalues pixelmask masklabels classlabels def getconfigself return maskformerconfig frombackboneanddecoderconfigs backboneconfigswinconfig depths1 1 1 1 embeddim16 hiddensize32 numheads1 1 2 2 decoderconfigdetrconfig decoderffndim64 decoderlayersself numhiddenlayers decoderattentionheadsself numattentionheads encoderffndim64 encoderlayersself numhiddenlayers encoderattentionheadsself numattentionheads numqueriesself numqueries dmodelself maskfeaturesize maskfeaturesizeself maskfeaturesize fpnfeaturesizeself maskfeaturesize numchannelsself numchannels numlabelsself numlabels def prepareconfigandinputsforcommonself config pixelvalues pixelmask self prepareconfigandinputs inputsdict pixelvalues pixelvalues pixelmask pixelmask return config inputsdict def checkoutputhiddenstateself output config encoderhiddenstates output encoderhiddenstates pixeldecoderhiddenstates output pixeldecoderhiddenstates transformerdecoderhiddenstates output transformerdecoderhiddenstates self parent asserttruelenencoderhiddenstates lenconfig backboneconfig depths self parent asserttruelenpixeldecoderhiddenstates lenconfig backboneconfig depths self parent asserttruelentransformerdecoderhiddenstates config decoderconfig decoderlayers def createandcheckmaskformermodelself config pixelvalues pixelmask outputhiddenstatesfalse with torch nograd model maskformermodelconfigconfig model totorchdevice model eval output modelpixelvaluespixelvalues pixelmaskpixelmask output modelpixelvalues outputhiddenstatestrue the correct shape of output transformerdecoderhiddenstates ensure the correcteness of the encoder and pixel decoder self parent assertequal output transformerdecoderlasthiddenstate shape self batchsize self numqueries self maskfeaturesize let s ensure the other two hidden state exists self parent asserttrueoutput pixeldecoderlasthiddenstate is not none self parent asserttrueoutput encoderlasthiddenstate is not none if outputhiddenstates self checkoutputhiddenstateoutput config def createandcheckmaskformerinstancesegmentationheadmodel self config pixelvalues pixelmask masklabels classlabels model maskformerforinstancesegmentationconfigconfig model totorchdevice model eval def commcheckonoutputresult let s still check that all the required stuff is there self parent asserttrueresult transformerdecoderlasthiddenstate is not none self parent asserttrueresult pixeldecoderlasthiddenstate is not none self parent asserttrueresult encoderlasthiddenstate 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 masksquerieslogits shape self batchsize self numqueries self minsize 4 self maxsize 4 1 for null class self parent assertequal result classquerieslogits shape self batchsize self numqueries self numlabels 1 with torch nograd result modelpixelvaluespixelvalues pixelmaskpixelmask result modelpixelvalues commcheckonoutputresult result model pixelvaluespixelvalues pixelmaskpixelmask masklabelsmasklabels classlabelsclasslabels commcheckonoutputresult self parent asserttrueresult loss is not none self parent assertequalresult loss shape torch size1 requiretorch class maskformermodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses maskformermodel maskformerforinstancesegmentation if istorchavailable else pipelinemodelmapping featureextraction maskformermodel imagesegmentation maskformerforinstancesegmentation if istorchavailable else isencoderdecoder false testpruning false testheadmasking false testmissingkeys false def setupself self modeltester maskformermodeltesterself self configtester configtesterself configclassmaskformerconfig hastextmodalityfalse def prepareforclassself inputsdict modelclass returnlabelsfalse inputsdict copy deepcopyinputsdict if returnlabels if modelclass in maskformerforinstancesegmentation inputsdictmasklabels torch zeros self modeltester batchsize self modeltester numlabels self modeltester minsize self modeltester maxsize dtypetorch float32 devicetorchdevice inputsdictclasslabels torch zeros self modeltester batchsize self modeltester numlabels dtypetorch long devicetorchdevice return inputsdict def testconfigself self configtester runcommontests def testmaskformermodelself config inputs self modeltester prepareconfigandinputsforcommon self modeltester createandcheckmaskformermodelconfig inputs outputhiddenstatesfalse def testmaskformerinstancesegmentationheadmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmaskformerinstancesegmentationheadmodelconfigandinputs unittest skipreasonmaskformer does not use inputsembeds def testinputsembedsself pass unittest skipreasonmaskformer does not have a getinputembeddings method def testmodelcommonattributesself pass unittest skipreasonmaskformer is not a generative model def testgeneratewithoutinputidsself pass unittest skipreasonmaskformer does not use token embeddings def testresizetokensembeddingsself pass requiretorchmultigpu unittest skip reasonmaskformer has some layers using addmodule which doesn t work well with nn dataparallel def testmultigpudataparallelforwardself pass slow def testmodelfrompretrainedself for modelname in facebookmaskformerswinsmallcoco model maskformermodel frompretrainedmodelname self assertisnotnonemodel def testmodelwithlabelsself size self modeltester minsize 2 inputs pixelvalues torch randn2 3 size devicetorchdevice masklabels torch randn2 10 size devicetorchdevice classlabels torch zeros2 10 devicetorchdevice long model maskformerforinstancesegmentationmaskformerconfig totorchdevice outputs modelinputs self asserttrueoutputs loss is not none def testhiddenstatesoutputself config inputs self modeltester prepareconfigandinputsforcommon self modeltester createandcheckmaskformermodelconfig inputs outputhiddenstatestrue def testattentionoutputsself config inputsdict self modeltester prepareconfigandinputsforcommon config returndict true for modelclass in self allmodelclasses inputsdictoutputattentions true inputsdictoutputhiddenstates false config returndict true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs attentions self assertequallenattentions self modeltester numhiddenlayers check that outputattentions also work using config del inputsdictoutputattentions config outputattentions true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs attentions self assertequallenattentions self modeltester numhiddenlayers outlen lenoutputs check attention is always last and order is fine inputsdictoutputattentions true inputsdictoutputhiddenstates true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass encoderhiddenstates pixeldecoderhiddenstates transformerdecoderhiddenstates hiddenstates addedhiddenstates 4 self assertequaloutlen addedhiddenstates lenoutputs selfattentions outputs attentions self assertequallenselfattentions self modeltester numhiddenlayers def testretaingradhiddenstatesattentionsself only maskformerforinstancesegmentation has the loss modelclass self allmodelclasses1 config pixelvalues pixelmask masklabels classlabels self modeltester prepareconfigandinputs config outputhiddenstates true config outputattentions true model modelclassconfig model totorchdevice model train outputs modelpixelvalues masklabelsmasklabels classlabelsclasslabels encoderhiddenstates outputs encoderhiddenstates0 encoderhiddenstates retaingrad pixeldecoderhiddenstates outputs pixeldecoderhiddenstates0 pixeldecoderhiddenstates retaingrad we requiresgradtrue in inputsembeds line 2152 the original implementation don t transformerdecoderhiddenstates outputs transformerdecoderhiddenstates0 transformerdecoderhiddenstates retaingrad attentions outputs attentions0 attentions retaingrad outputs loss backwardretaingraphtrue self assertisnotnoneencoderhiddenstates grad self assertisnotnonepixeldecoderhiddenstates grad self assertisnotnonetransformerdecoderhiddenstates grad self assertisnotnoneattentions grad tolerance 1e4 we will verify our results on an image of cute cats def prepareimg image image open testsfixturestestssamplescoco000000039769 png return image requirevision slow class maskformermodelintegrationtestunittest testcase cachedproperty def defaultimageprocessorself return maskformerimageprocessor frompretrainedfacebookmaskformerswinsmallcoco if isvisionavailable else none def testinferencenoheadself model maskformermodel frompretrainedfacebookmaskformerswinsmallcoco totorchdevice imageprocessor self defaultimageprocessor image prepareimg inputs imageprocessorimage returntensorspt totorchdevice inputsshape inputspixelvalues shape check size is divisible by 32 self asserttrueinputsshape1 32 0 and inputsshape2 32 0 check size self assertequalinputsshape 1 3 800 1088 with torch nograd outputs modelinputs expectedslicehiddenstate torch tensor 0 0482 0 9228 0 4951 0 2547 0 8017 0 8527 0 0069 0 3385 0 0089 totorchdevice self asserttrue torch allclose outputs encoderlasthiddenstate0 0 3 3 expectedslicehiddenstate atoltolerance expectedslicehiddenstate torch tensor 0 8422 0 8434 0 9718 1 0144 0 5565 0 4195 1 0038 0 4484 0 1961 totorchdevice self asserttrue torch allclose outputs pixeldecoderlasthiddenstate0 0 3 3 expectedslicehiddenstate atoltolerance expectedslicehiddenstate torch tensor 0 2852 0 0159 0 9735 0 6254 0 1858 0 8529 0 0680 0 4116 1 8413 totorchdevice self asserttrue torch allclose outputs transformerdecoderlasthiddenstate0 3 3 expectedslicehiddenstate atoltolerance def testinferenceinstancesegmentationheadself model maskformerforinstancesegmentation frompretrainedfacebookmaskformerswinsmallcoco totorchdevice eval imageprocessor self defaultimageprocessor image prepareimg inputs imageprocessorimage returntensorspt totorchdevice inputsshape inputspixelvalues shape check size is divisible by 32 self asserttrueinputsshape1 32 0 and inputsshape2 32 0 check size self assertequalinputsshape 1 3 800 1088 with torch nograd outputs modelinputs masksquerieslogits masksquerieslogits outputs masksquerieslogits self assertequal masksquerieslogits shape 1 model config decoderconfig numqueries inputsshape2 4 inputsshape1 4 expectedslice 1 3737124 1 7724937 1 9364233 1 5977281 1 9867939 2 1523695 1 5795398 1 9269832 2 093942 expectedslice torch tensorexpectedslice totorchdevice self asserttruetorch allclosemasksquerieslogits0 0 3 3 expectedslice atoltolerance classquerieslogits classquerieslogits outputs classquerieslogits self assertequal classquerieslogits shape 1 model config decoderconfig numqueries model config numlabels 1 expectedslice torch tensor 1 6512e00 5 2572e00 3 3519e00 3 6169e02 5 9025e00 2 9313e00 1 0766e04 7 7630e00 5 1263e00 totorchdevice self asserttruetorch allcloseoutputs classquerieslogits0 3 3 expectedslice atoltolerance def testinferenceinstancesegmentationheadresnetbackboneself model maskformerforinstancesegmentation frompretrainedfacebookmaskformerresnet101cocostuff totorchdevice eval imageprocessor self defaultimageprocessor image prepareimg inputs imageprocessorimage returntensorspt totorchdevice inputsshape inputspixelvalues shape check size is divisible by 32 self asserttrueinputsshape1 32 0 and inputsshape2 32 0 check size self assertequalinputsshape 1 3 800 1088 with torch nograd outputs modelinputs masksquerieslogits masksquerieslogits outputs masksquerieslogits self assertequal masksquerieslogits shape 1 model config decoderconfig numqueries inputsshape2 4 inputsshape1 4 expectedslice 0 9046 2 6366 4 6062 3 4179 5 7890 8 8057 4 9179 7 6560 10 7711 expectedslice torch tensorexpectedslice totorchdevice self asserttruetorch allclosemasksquerieslogits0 0 3 3 expectedslice atoltolerance classquerieslogits classquerieslogits outputs classquerieslogits self assertequal classquerieslogits shape 1 model config decoderconfig numqueries model config numlabels 1 expectedslice torch tensor 4 7188 3 2585 2 8857 6 6871 2 9181 1 2487 7 2449 2 2764 2 1874 totorchdevice self asserttruetorch allcloseoutputs classquerieslogits0 3 3 expectedslice atoltolerance requiretorchaccelerator requiretorchfp16 def testinferencefp16self model maskformerforinstancesegmentation frompretrainedfacebookmaskformerresnet101cocostuff totorchdevice dtypetorch float16 eval imageprocessor self defaultimageprocessor image prepareimg inputs imageprocessorimage returntensorspt totorchdevice dtypetorch float16 with torch nograd modelinputs def testwithsegmentationmapsandlossself model maskformerforinstancesegmentation frompretrainedfacebookmaskformerswinsmallcoco totorchdevice eval imageprocessor self defaultimageprocessor inputs imageprocessor np zeros3 400 333 np zeros3 400 333 segmentationmapsnp zeros384 384 astypenp float32 np zeros384 384 astypenp float32 returntensorspt inputspixelvalues inputspixelvalues totorchdevice inputsmasklabels el totorchdevice for el in inputsmasklabels inputsclasslabels el totorchdevice for el in inputsclasslabels with torch nograd outputs modelinputs self asserttrueoutputs loss is not none coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch maskformer model this is passed to the decoder config we add it to the model tester here for testing the correct shape of output transformer_decoder_hidden_states ensure the correcteness of the encoder and pixel decoder let s ensure the other two hidden state exists let s still check that all the required stuff is there okay now we need to check the logits shape due to the encoder compression masks have a 4 spatial size 1 for null class check that output_attentions also work using config check attention is always last and order is fine encoder_hidden_states pixel_decoder_hidden_states transformer_decoder_hidden_states hidden_states only maskformerforinstancesegmentation has the loss we requires_grad true in inputs_embeds line 2152 the original implementation don t we will verify our results on an image of cute cats check size is divisible by 32 check size check size is divisible by 32 check size masks_queries_logits class_queries_logits check size is divisible by 32 check size masks_queries_logits class_queries_logits
import copy
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_accelerator,
    require_torch_fp16,
    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 MaskFormerModelTester:
    def __init__(
        self,
        parent,
        batch_size=2,
        is_training=True,
        use_auxiliary_loss=False,
        num_queries=10,
        num_channels=3,
        min_size=32 * 4,
        max_size=32 * 6,
        num_labels=4,
        mask_feature_size=32,
        num_hidden_layers=2,
        num_attention_heads=2,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.is_training = is_training
        self.use_auxiliary_loss = use_auxiliary_loss
        self.num_queries = num_queries
        self.num_channels = num_channels
        self.min_size = min_size
        self.max_size = max_size
        self.num_labels = num_labels
        self.mask_feature_size = mask_feature_size
        
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads

    def prepare_config_and_inputs(self):
        pixel_values = floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]).to(
            torch_device
        )

        pixel_mask = torch.ones([self.batch_size, self.min_size, self.max_size], device=torch_device)

        mask_labels = (
            torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size], device=torch_device) > 0.5
        ).float()
        class_labels = (torch.rand((self.batch_size, self.num_labels), device=torch_device) > 0.5).long()

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

    def get_config(self):
        return MaskFormerConfig.from_backbone_and_decoder_configs(
            backbone_config=SwinConfig(
                depths=[1, 1, 1, 1],
                embed_dim=16,
                hidden_size=32,
                num_heads=[1, 1, 2, 2],
            ),
            decoder_config=DetrConfig(
                decoder_ffn_dim=64,
                decoder_layers=self.num_hidden_layers,
                decoder_attention_heads=self.num_attention_heads,
                encoder_ffn_dim=64,
                encoder_layers=self.num_hidden_layers,
                encoder_attention_heads=self.num_attention_heads,
                num_queries=self.num_queries,
                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 prepare_config_and_inputs_for_common(self):
        config, pixel_values, pixel_mask, _, _ = self.prepare_config_and_inputs()
        inputs_dict = {"pixel_values": pixel_values, "pixel_mask": pixel_mask}
        return config, inputs_dict

    def check_output_hidden_state(self, output, config):
        encoder_hidden_states = output.encoder_hidden_states
        pixel_decoder_hidden_states = output.pixel_decoder_hidden_states
        transformer_decoder_hidden_states = output.transformer_decoder_hidden_states

        self.parent.assertTrue(len(encoder_hidden_states), len(config.backbone_config.depths))
        self.parent.assertTrue(len(pixel_decoder_hidden_states), len(config.backbone_config.depths))
        self.parent.assertTrue(len(transformer_decoder_hidden_states), config.decoder_config.decoder_layers)

    def create_and_check_maskformer_model(self, config, pixel_values, pixel_mask, output_hidden_states=False):
        with torch.no_grad():
            model = MaskFormerModel(config=config)
            model.to(torch_device)
            model.eval()

            output = model(pixel_values=pixel_values, pixel_mask=pixel_mask)
            output = model(pixel_values, output_hidden_states=True)
        
        
        self.parent.assertEqual(
            output.transformer_decoder_last_hidden_state.shape,
            (self.batch_size, self.num_queries, self.mask_feature_size),
        )
        
        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(output, config)

    def create_and_check_maskformer_instance_segmentation_head_model(
        self, config, pixel_values, pixel_mask, mask_labels, class_labels
    ):
        model = MaskFormerForInstanceSegmentation(config=config)
        model.to(torch_device)
        model.eval()

        def comm_check_on_output(result):
            
            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)
            
            
            self.parent.assertEqual(
                result.masks_queries_logits.shape,
                (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4),
            )
            
            self.parent.assertEqual(
                result.class_queries_logits.shape, (self.batch_size, self.num_queries, self.num_labels + 1)
            )

        with torch.no_grad():
            result = model(pixel_values=pixel_values, pixel_mask=pixel_mask)
            result = model(pixel_values)

            comm_check_on_output(result)

            result = model(
                pixel_values=pixel_values, pixel_mask=pixel_mask, mask_labels=mask_labels, class_labels=class_labels
            )

        comm_check_on_output(result)

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


@require_torch
class MaskFormerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
    pipeline_model_mapping = (
        {"feature-extraction": MaskFormerModel, "image-segmentation": MaskFormerForInstanceSegmentation}
        if is_torch_available()
        else {}
    )

    is_encoder_decoder = False
    test_pruning = False
    test_head_masking = False
    test_missing_keys = False

    def setUp(self):
        self.model_tester = MaskFormerModelTester(self)
        self.config_tester = ConfigTester(self, config_class=MaskFormerConfig, has_text_modality=False)

    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
        inputs_dict = copy.deepcopy(inputs_dict)

        if return_labels:
            if model_class in [MaskFormerForInstanceSegmentation]:
                inputs_dict["mask_labels"] = torch.zeros(
                    (
                        self.model_tester.batch_size,
                        self.model_tester.num_labels,
                        self.model_tester.min_size,
                        self.model_tester.max_size,
                    ),
                    dtype=torch.float32,
                    device=torch_device,
                )
                inputs_dict["class_labels"] = torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.num_labels), dtype=torch.long, device=torch_device
                )

        return inputs_dict

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

    def test_maskformer_model(self):
        config, inputs = self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.create_and_check_maskformer_model(config, **inputs, output_hidden_states=False)

    def test_maskformer_instance_segmentation_head_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*config_and_inputs)

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

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

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

    @unittest.skip(reason="MaskFormer does not use token embeddings")
    def test_resize_tokens_embeddings(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 test_multi_gpu_data_parallel_forward(self):
        pass

    @slow
    def test_model_from_pretrained(self):
        for model_name in ["facebook/maskformer-swin-small-coco"]:
            model = MaskFormerModel.from_pretrained(model_name)
            self.assertIsNotNone(model)

    def test_model_with_labels(self):
        size = (self.model_tester.min_size,) * 2
        inputs = {
            "pixel_values": torch.randn((2, 3, *size), device=torch_device),
            "mask_labels": torch.randn((2, 10, *size), device=torch_device),
            "class_labels": torch.zeros(2, 10, device=torch_device).long(),
        }

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

    def test_hidden_states_output(self):
        config, inputs = self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.create_and_check_maskformer_model(config, **inputs, output_hidden_states=True)

    def test_attention_outputs(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.return_dict = True

        for model_class in self.all_model_classes:
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = False
            config.return_dict = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

            
            del inputs_dict["output_attentions"]
            config.output_attentions = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)
            out_len = len(outputs)

            
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()
            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            
            added_hidden_states = 4
            self.assertEqual(out_len + added_hidden_states, len(outputs))

            self_attentions = outputs.attentions
            self.assertEqual(len(self_attentions), self.model_tester.num_hidden_layers)

    def test_retain_grad_hidden_states_attentions(self):
        
        model_class = self.all_model_classes[1]
        config, pixel_values, pixel_mask, mask_labels, class_labels = self.model_tester.prepare_config_and_inputs()
        config.output_hidden_states = True
        config.output_attentions = True

        model = model_class(config)
        model.to(torch_device)
        model.train()

        outputs = model(pixel_values, mask_labels=mask_labels, class_labels=class_labels)

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

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

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

        outputs.loss.backward(retain_graph=True)

        self.assertIsNotNone(encoder_hidden_states.grad)
        self.assertIsNotNone(pixel_decoder_hidden_states.grad)
        self.assertIsNotNone(transformer_decoder_hidden_states.grad)
        self.assertIsNotNone(attentions.grad)


TOLERANCE = 1e-4



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


@require_vision
@slow
class MaskFormerModelIntegrationTest(unittest.TestCase):
    @cached_property
    def default_image_processor(self):
        return (
            MaskFormerImageProcessor.from_pretrained("facebook/maskformer-swin-small-coco")
            if is_vision_available()
            else None
        )

    def test_inference_no_head(self):
        model = MaskFormerModel.from_pretrained("facebook/maskformer-swin-small-coco").to(torch_device)
        image_processor = self.default_image_processor
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").to(torch_device)
        inputs_shape = inputs["pixel_values"].shape
        
        self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0)
        
        self.assertEqual(inputs_shape, (1, 3, 800, 1088))

        with torch.no_grad():
            outputs = model(**inputs)

        expected_slice_hidden_state = torch.tensor(
            [[-0.0482, 0.9228, 0.4951], [-0.2547, 0.8017, 0.8527], [-0.0069, 0.3385, -0.0089]]
        ).to(torch_device)
        self.assertTrue(
            torch.allclose(
                outputs.encoder_last_hidden_state[0, 0, :3, :3], expected_slice_hidden_state, atol=TOLERANCE
            )
        )

        expected_slice_hidden_state = torch.tensor(
            [[-0.8422, -0.8434, -0.9718], [-1.0144, -0.5565, -0.4195], [-1.0038, -0.4484, -0.1961]]
        ).to(torch_device)
        self.assertTrue(
            torch.allclose(
                outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3], expected_slice_hidden_state, atol=TOLERANCE
            )
        )

        expected_slice_hidden_state = torch.tensor(
            [[0.2852, -0.0159, 0.9735], [0.6254, 0.1858, 0.8529], [-0.0680, -0.4116, 1.8413]]
        ).to(torch_device)
        self.assertTrue(
            torch.allclose(
                outputs.transformer_decoder_last_hidden_state[0, :3, :3], expected_slice_hidden_state, atol=TOLERANCE
            )
        )

    def test_inference_instance_segmentation_head(self):
        model = (
            MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-small-coco")
            .to(torch_device)
            .eval()
        )
        image_processor = self.default_image_processor
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").to(torch_device)
        inputs_shape = inputs["pixel_values"].shape
        
        self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0)
        
        self.assertEqual(inputs_shape, (1, 3, 800, 1088))

        with torch.no_grad():
            outputs = model(**inputs)
        
        masks_queries_logits = 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),
        )
        expected_slice = [
            [-1.3737124, -1.7724937, -1.9364233],
            [-1.5977281, -1.9867939, -2.1523695],
            [-1.5795398, -1.9269832, -2.093942],
        ]
        expected_slice = torch.tensor(expected_slice).to(torch_device)
        self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3], expected_slice, atol=TOLERANCE))
        
        class_queries_logits = outputs.class_queries_logits
        self.assertEqual(
            class_queries_logits.shape, (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)
        )
        expected_slice = torch.tensor(
            [
                [1.6512e00, -5.2572e00, -3.3519e00],
                [3.6169e-02, -5.9025e00, -2.9313e00],
                [1.0766e-04, -7.7630e00, -5.1263e00],
            ]
        ).to(torch_device)
        self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3], expected_slice, atol=TOLERANCE))

    def test_inference_instance_segmentation_head_resnet_backbone(self):
        model = (
            MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-resnet101-coco-stuff")
            .to(torch_device)
            .eval()
        )
        image_processor = self.default_image_processor
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").to(torch_device)
        inputs_shape = inputs["pixel_values"].shape
        
        self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0)
        
        self.assertEqual(inputs_shape, (1, 3, 800, 1088))

        with torch.no_grad():
            outputs = model(**inputs)
        
        masks_queries_logits = 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),
        )
        expected_slice = [[-0.9046, -2.6366, -4.6062], [-3.4179, -5.7890, -8.8057], [-4.9179, -7.6560, -10.7711]]
        expected_slice = torch.tensor(expected_slice).to(torch_device)
        self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3], expected_slice, atol=TOLERANCE))
        
        class_queries_logits = outputs.class_queries_logits
        self.assertEqual(
            class_queries_logits.shape, (1, model.config.decoder_config.num_queries, model.config.num_labels + 1)
        )
        expected_slice = torch.tensor(
            [[4.7188, -3.2585, -2.8857], [6.6871, -2.9181, -1.2487], [7.2449, -2.2764, -2.1874]]
        ).to(torch_device)
        self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3], expected_slice, atol=TOLERANCE))

    @require_torch_accelerator
    @require_torch_fp16
    def test_inference_fp16(self):
        model = (
            MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-resnet101-coco-stuff")
            .to(torch_device, dtype=torch.float16)
            .eval()
        )
        image_processor = self.default_image_processor
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").to(torch_device, dtype=torch.float16)

        with torch.no_grad():
            _ = model(**inputs)

    def test_with_segmentation_maps_and_loss(self):
        model = (
            MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-small-coco")
            .to(torch_device)
            .eval()
        )
        image_processor = self.default_image_processor

        inputs = image_processor(
            [np.zeros((3, 400, 333)), np.zeros((3, 400, 333))],
            segmentation_maps=[np.zeros((384, 384)).astype(np.float32), np.zeros((384, 384)).astype(np.float32)],
            return_tensors="pt",
        )

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

        with torch.no_grad():
            outputs = model(**inputs)

        self.assertTrue(outputs.loss is not None)