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 oneformer model import copy import inspect import unittest import numpy as np from tests testmodelingcommon import floatstensor from transformers import oneformerconfig 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 oneformerforuniversalsegmentation oneformermodel if isvisionavailable from transformers import oneformerprocessor if isvisionavailable from pil import image def configzeroinitconfig configsnoinit copy deepcopyconfig for key in configsnoinit dict keys if range in key or std in key or initializerfactor in key or layerscale in key setattrconfigsnoinit key 1e10 return configsnoinit class oneformermodeltester def init self parent batchsize2 istrainingtrue vocabsize99 useauxiliarylossfalse numqueries10 numchannels3 minsize32 8 maxsize32 8 numlabels4 hiddendim64 sequencelength77 nctx4 self parent parent self batchsize batchsize self istraining istraining self vocabsize vocabsize self useauxiliaryloss useauxiliaryloss self numqueries numqueries self numchannels numchannels self minsize minsize self maxsize maxsize self numlabels numlabels self hiddendim hiddendim self sequencelength sequencelength self nctx nctx def prepareconfigandinputsself pixelvalues floatstensorself batchsize self numchannels self minsize self maxsize to torchdevice taskinputs torch randinthighself vocabsize sizeself batchsize self sequencelength totorchdevice long pixelmask torch onesself batchsize self minsize self maxsize devicetorchdevice textinputs torch randint highself vocabsize sizeself batchsize self numqueries self nctx self sequencelength totorchdevice long 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 taskinputs textinputs pixelmask masklabels classlabels def getconfigself config oneformerconfig textencodervocabsizeself vocabsize hiddensizeself hiddendim numqueriesself numqueries numlabelsself numlabels encoderfeedforwarddim32 dimfeedforward64 encoderlayers2 decoderlayers2 config backboneconfig embeddim 16 config backboneconfig depths 1 1 1 1 config backboneconfig hiddensize 16 config backboneconfig numchannels self numchannels config backboneconfig numheads 1 1 2 2 config hiddendim self hiddendim config maskdim self hiddendim config convdim self hiddendim config textencoderwidth self hiddendim config taskseqlen self sequencelength config maxseqlen self sequencelength config textencodercontextlength self sequencelength config textencodernctx self nctx return config def prepareconfigandinputsforcommonself config pixelvalues taskinputs pixelmask self prepareconfigandinputs inputsdict pixelvalues pixelvalues pixelmask pixelmask taskinputs taskinputs 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 config encoderlayers self parent asserttruelentransformerdecoderhiddenstates config decoderlayers 1 def createandcheckoneformermodel self config pixelvalues taskinputs pixelmask outputhiddenstatesfalse with torch nograd model oneformermodelconfigconfig model totorchdevice model eval output modelpixelvaluespixelvalues taskinputstaskinputs pixelmaskpixelmask output modelpixelvalues taskinputstaskinputs outputhiddenstatestrue the correct shape of output transformerdecoderhiddenstates ensure the correcteness of the encoder and pixel decoder self parent assertequal output transformerdecoderobjectqueries shape self batchsize self numqueries self hiddendim let s ensure the other two hidden state exists self parent asserttrueoutput pixeldecoderhiddenstates is not none self parent asserttrueoutput encoderhiddenstates is not none if outputhiddenstates self checkoutputhiddenstateoutput config def createandcheckoneformeruniversalsegmentationheadmodel self config pixelvalues taskinputs textinputs pixelmask masklabels classlabels model oneformerforuniversalsegmentationconfigconfig model totorchdevice model eval def commcheckonoutputresult let s still check that all the required stuff is there self parent asserttrueresult transformerdecoderhiddenstates is not none self parent asserttrueresult pixeldecoderhiddenstates is not none self parent asserttrueresult encoderhiddenstates 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 taskinputstaskinputs pixelmaskpixelmask result modelpixelvalues taskinputs commcheckonoutputresult config istraining true model oneformerforuniversalsegmentationconfigconfig model totorchdevice model eval with torch nograd result model pixelvaluespixelvalues taskinputstaskinputs pixelmaskpixelmask masklabelsmasklabels classlabelsclasslabels textinputstextinputs commcheckonoutputresult self parent asserttrueresult loss is not none self parent assertequalresult loss shape torch size1 requiretorch class oneformermodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses oneformermodel oneformerforuniversalsegmentation if istorchavailable else pipelinemodelmapping featureextraction oneformermodel if istorchavailable else isencoderdecoder false testpruning false testheadmasking false testmissingkeys false todo fix the failed tests when this model gets more usage def ispipelinetesttoskip self pipelinetestcassename configclass modelarchitecture tokenizername processorname if pipelinetestcassename featureextractionpipelinetests return true return false def setupself self modeltester oneformermodeltesterself self configtester configtesterself configclassoneformerconfig hastextmodalityfalse def testconfigself self configtester runcommontests def testoneformermodelself config inputs self modeltester prepareconfigandinputsforcommon self modeltester createandcheckoneformermodelconfig inputs outputhiddenstatesfalse def testoneformeruniversalsegmentationheadmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckoneformeruniversalsegmentationheadmodelconfigandinputs def testmodelmaininputnameself for modelclass in self allmodelclasses modelsignature inspect signaturegetattrmodelclass forward the main input is the name of the argument after self observedmaininputname listmodelsignature parameters keys1 3 self assertequalmodelclass maininputname observedmaininputname unittest skipreasononeformer uses two main inputs def testtorchscriptsimpleself pass unittest skipreasononeformer uses two main inputs def testtorchscriptoutputattentionsself pass unittest skipreasononeformer uses two main inputs def testtorchscriptoutputhiddenstateself pass unittest skipreasononeformer does not use inputsembeds def testinputsembedsself pass unittest skipreasononeformer does not have a getinputembeddings method def testmodelcommonattributesself pass unittest skipreasononeformer is not a generative model def testgeneratewithoutinputidsself pass unittest skipreasononeformer does not use token embeddings def testresizetokensembeddingsself pass requiretorchmultigpu unittest skip reasononeformer has some layers using addmodule which doesn t work well with nn dataparallel def testmultigpudataparallelforwardself pass def testforwardsignatureself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig signature inspect signaturemodel forward signature parameters is an ordereddict so argnames order is deterministic argnames signature parameters keys expectedargnames pixelvalues taskinputs self assertlistequalargnames 2 expectedargnames slow def testmodelfrompretrainedself for modelname in shilabsoneformerade20kswintiny model oneformermodel frompretrainedmodelname self assertisnotnonemodel def testmodelwithlabelsself size self modeltester minsize 2 inputs pixelvalues torch randn2 3 size devicetorchdevice taskinputs torch randinthighself modeltester vocabsize size2 77 devicetorchdevice long textinputs torch randint highself modeltester vocabsize size2 6 77 devicetorchdevice long masklabels torch randn2 150 size devicetorchdevice classlabels torch zeros2 150 devicetorchdevice long config self modeltester getconfig config istraining true model oneformerforuniversalsegmentationconfig totorchdevice outputs modelinputs self asserttrueoutputs loss is not none def testhiddenstatesoutputself config inputs self modeltester prepareconfigandinputsforcommon self modeltester createandcheckoneformermodelconfig inputs outputhiddenstatestrue def testattentionoutputsself config inputs self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig totorchdevice outputs modelinputs outputattentionstrue self asserttrueoutputs attentions is not none def testinitializationself config inputsdict self modeltester prepareconfigandinputsforcommon config contrastivetemperature 1 configsnoinit configzeroinitconfig for modelclass in self allmodelclasses model modelclassconfigconfigsnoinit for name param in model namedparameters if param requiresgrad self assertin param data mean 1e9 round 1e9 item 0 0 1 0 msgfparameter name of model modelclass seems not properly initialized def testtrainingself if not self modeltester istraining return only oneformerforuniversalsegmentation has the loss modelclass self allmodelclasses1 config pixelvalues taskinputs textinputs pixelmask masklabels classlabels self modeltester prepareconfigandinputs config istraining true model modelclassconfig model totorchdevice model train loss model pixelvalues taskinputs textinputstextinputs masklabelsmasklabels classlabelsclasslabels loss loss backward def testretaingradhiddenstatesattentionsself only oneformerforuniversalsegmentation has the loss modelclass self allmodelclasses1 config pixelvalues taskinputs textinputs pixelmask masklabels classlabels self modeltester prepareconfigandinputs config outputhiddenstates true config outputattentions true config istraining true model modelclassconfig model totorchdevice model train outputs model pixelvalues taskinputs textinputstextinputs masklabelsmasklabels classlabelsclasslabels encoderhiddenstates outputs encoderhiddenstates0 encoderhiddenstates retaingrad pixeldecoderhiddenstates outputs pixeldecoderhiddenstates0 pixeldecoderhiddenstates retaingrad transformerdecoderclasspredictions outputs transformerdecoderclasspredictions transformerdecoderclasspredictions retaingrad transformerdecodermaskpredictions outputs transformerdecodermaskpredictions transformerdecodermaskpredictions retaingrad attentions outputs attentions00 attentions retaingrad outputs loss backwardretaingraphtrue self assertisnotnoneencoderhiddenstates grad self assertisnotnonepixeldecoderhiddenstates grad self assertisnotnonetransformerdecoderclasspredictions grad self assertisnotnonetransformerdecodermaskpredictions 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 oneformermodelintegrationtestunittest testcase cachedproperty def modelcheckpointsself return shilabsoneformerade20kswintiny cachedproperty def defaultprocessorself return oneformerprocessor frompretrainedself modelcheckpoints if isvisionavailable else none def testinferencenoheadself model oneformermodel frompretrainedself modelcheckpoints totorchdevice processor self defaultprocessor image prepareimg inputs processorimage semantic returntensorspt totorchdevice inputsshape inputspixelvalues shape check size self assertequalinputsshape 1 3 512 682 taskinputsshape inputstaskinputs shape check size self assertequaltaskinputsshape 1 77 with torch nograd outputs modelinputs expectedslicehiddenstate torch tensor 0 2723 0 8280 0 6026 1 2699 1 1257 1 1444 1 1344 0 6153 0 4177 totorchdevice self asserttrue torch allclose outputs encoderhiddenstates10 0 3 3 expectedslicehiddenstate atoltolerance expectedslicehiddenstate torch tensor 1 0581 1 2276 1 2003 1 1903 1 2925 1 2862 1 158 1 2559 1 3216 totorchdevice self asserttrue torch allclose outputs pixeldecoderhiddenstates00 0 3 3 expectedslicehiddenstate atoltolerance expectedslicehiddenstate torch tensor 3 0668 1 1833 5 1103 3 344 3 362 5 1101 2 6017 4 3613 4 1444 totorchdevice self asserttrue torch allclose outputs transformerdecoderclasspredictions0 3 3 expectedslicehiddenstate atoltolerance def testinferenceuniversalsegmentationheadself model oneformerforuniversalsegmentation frompretrainedself modelcheckpoints totorchdevice eval processor self defaultprocessor image prepareimg inputs processorimage semantic returntensorspt totorchdevice inputsshape inputspixelvalues shape check size self assertequalinputsshape 1 3 512 682 with torch nograd outputs modelinputs masksquerieslogits masksquerieslogits outputs masksquerieslogits self assertequal masksquerieslogits shape 1 model config numqueries inputsshape2 4 inputsshape1 2 4 expectedslice 3 1848 4 2141 4 1993 2 9000 3 5721 3 6603 2 5358 3 0883 3 6168 expectedslice torch tensorexpectedslice totorchdevice self asserttruetorch allclosemasksquerieslogits0 0 3 3 expectedslice atoltolerance classquerieslogits classquerieslogits outputs classquerieslogits self assertequal classquerieslogits shape 1 model config numqueries model config numlabels 1 expectedslice torch tensor 3 0668 1 1833 5 1103 3 3440 3 3620 5 1101 2 6017 4 3613 4 1444 totorchdevice self asserttruetorch allcloseclassquerieslogits0 3 3 expectedslice atoltolerance requiretorchaccelerator requiretorchfp16 def testinferencefp16self model oneformerforuniversalsegmentation frompretrainedself modelcheckpoints totorchdevice dtypetorch float16 eval processor self defaultprocessor image prepareimg inputs processorimage semantic returntensorspt totorchdevice dtypetorch float16 with torch nograd modelinputs def testwithsegmentationmapsandlossself dummymodel oneformerforuniversalsegmentation frompretrainedself modelcheckpoints processor self defaultprocessor processor imageprocessor numtext dummymodel config numqueries dummymodel config textencodernctx dummymodel config istraining true model oneformerforuniversalsegmentationdummymodel config totorchdevice eval del dummymodel inputs processor np zeros3 512 640 np zeros3 512 640 semantic semantic segmentationmapsnp zeros384 384 astypenp float32 np zeros384 384 astypenp float32 returntensorspt inputspixelvalues inputspixelvalues totorchdevice inputstaskinputs inputstaskinputs totorchdevice inputstextinputs inputstextinputs 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 oneformer model 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 todo fix the failed tests when this model gets more usage the main input is the name of the argument after self signature parameters is an ordereddict so arg_names order is deterministic only oneformerforuniversalsegmentation has the loss only oneformerforuniversalsegmentation has the loss we will verify our results on an image of cute cats check size check size check size masks_queries_logits class_queries_logits
import copy
import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import OneFormerConfig, 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 OneFormerForUniversalSegmentation, OneFormerModel

    if is_vision_available():
        from transformers import OneFormerProcessor

if is_vision_available():
    from PIL import Image


def _config_zero_init(config):
    configs_no_init = copy.deepcopy(config)
    for key in configs_no_init.__dict__.keys():
        if "_range" in key or "_std" in key or "initializer_factor" in key or "layer_scale" in key:
            setattr(configs_no_init, key, 1e-10)
    return configs_no_init


class OneFormerModelTester:
    def __init__(
        self,
        parent,
        batch_size=2,
        is_training=True,
        vocab_size=99,
        use_auxiliary_loss=False,
        num_queries=10,
        num_channels=3,
        min_size=32 * 8,
        max_size=32 * 8,
        num_labels=4,
        hidden_dim=64,
        sequence_length=77,
        n_ctx=4,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.is_training = is_training
        self.vocab_size = vocab_size
        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.hidden_dim = hidden_dim
        self.sequence_length = sequence_length
        self.n_ctx = n_ctx

    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
        )

        task_inputs = (
            torch.randint(high=self.vocab_size, size=(self.batch_size, self.sequence_length)).to(torch_device).long()
        )

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

        text_inputs = (
            torch.randint(
                high=self.vocab_size, size=(self.batch_size, self.num_queries - self.n_ctx, self.sequence_length)
            )
            .to(torch_device)
            .long()
        )

        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, task_inputs, text_inputs, pixel_mask, mask_labels, class_labels

    def get_config(self):
        config = OneFormerConfig(
            text_encoder_vocab_size=self.vocab_size,
            hidden_size=self.hidden_dim,
            num_queries=self.num_queries,
            num_labels=self.num_labels,
            encoder_feedforward_dim=32,
            dim_feedforward=64,
            encoder_layers=2,
            decoder_layers=2,
        )

        config.backbone_config.embed_dim = 16
        config.backbone_config.depths = [1, 1, 1, 1]
        config.backbone_config.hidden_size = 16
        config.backbone_config.num_channels = self.num_channels
        config.backbone_config.num_heads = [1, 1, 2, 2]

        config.hidden_dim = self.hidden_dim
        config.mask_dim = self.hidden_dim
        config.conv_dim = self.hidden_dim

        config.text_encoder_width = self.hidden_dim
        config.task_seq_len = self.sequence_length
        config.max_seq_len = self.sequence_length
        config.text_encoder_context_length = self.sequence_length
        config.text_encoder_n_ctx = self.n_ctx

        return config

    def prepare_config_and_inputs_for_common(self):
        config, pixel_values, task_inputs, pixel_mask, _, _, _ = self.prepare_config_and_inputs()
        inputs_dict = {"pixel_values": pixel_values, "pixel_mask": pixel_mask, "task_inputs": task_inputs}
        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), config.encoder_layers)
        self.parent.assertTrue(len(transformer_decoder_hidden_states), config.decoder_layers - 1)

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

            output = model(pixel_values=pixel_values, task_inputs=task_inputs, pixel_mask=pixel_mask)
            output = model(pixel_values, task_inputs=task_inputs, output_hidden_states=True)
        
        
        self.parent.assertEqual(
            output.transformer_decoder_object_queries.shape,
            (self.batch_size, self.num_queries, self.hidden_dim),
        )
        
        self.parent.assertTrue(output.pixel_decoder_hidden_states is not None)
        self.parent.assertTrue(output.encoder_hidden_states is not None)

        if output_hidden_states:
            self.check_output_hidden_state(output, config)

    def create_and_check_oneformer_universal_segmentation_head_model(
        self, config, pixel_values, task_inputs, text_inputs, pixel_mask, mask_labels, class_labels
    ):
        model = OneFormerForUniversalSegmentation(config=config)
        model.to(torch_device)
        model.eval()

        def comm_check_on_output(result):
            
            self.parent.assertTrue(result.transformer_decoder_hidden_states is not None)
            self.parent.assertTrue(result.pixel_decoder_hidden_states is not None)
            self.parent.assertTrue(result.encoder_hidden_states 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, task_inputs=task_inputs, pixel_mask=pixel_mask)
            result = model(pixel_values, task_inputs)

            comm_check_on_output(result)

        config.is_training = True
        model = OneFormerForUniversalSegmentation(config=config)
        model.to(torch_device)
        model.eval()

        with torch.no_grad():
            result = model(
                pixel_values=pixel_values,
                task_inputs=task_inputs,
                pixel_mask=pixel_mask,
                mask_labels=mask_labels,
                class_labels=class_labels,
                text_inputs=text_inputs,
            )

        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 OneFormerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (OneFormerModel, OneFormerForUniversalSegmentation) if is_torch_available() else ()
    pipeline_model_mapping = {"feature-extraction": OneFormerModel} if is_torch_available() else {}

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

    
    def is_pipeline_test_to_skip(
        self, pipeline_test_casse_name, config_class, model_architecture, tokenizer_name, processor_name
    ):
        if pipeline_test_casse_name == "FeatureExtractionPipelineTests":
            return True

        return False

    def setUp(self):
        self.model_tester = OneFormerModelTester(self)
        self.config_tester = ConfigTester(self, config_class=OneFormerConfig, has_text_modality=False)

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

    def test_oneformer_model(self):
        config, inputs = self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.create_and_check_oneformer_model(config, **inputs, output_hidden_states=False)

    def test_oneformer_universal_segmentation_head_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_oneformer_universal_segmentation_head_model(*config_and_inputs)

    def test_model_main_input_name(self):
        for model_class in self.all_model_classes:
            model_signature = inspect.signature(getattr(model_class, "forward"))
            
            observed_main_input_name = list(model_signature.parameters.keys())[1:3]
            self.assertEqual(model_class.main_input_name, observed_main_input_name)

    @unittest.skip(reason="OneFormer uses two main inputs")
    def test_torchscript_simple(self):
        pass

    @unittest.skip(reason="OneFormer uses two main inputs")
    def test_torchscript_output_attentions(self):
        pass

    @unittest.skip(reason="OneFormer uses two main inputs")
    def test_torchscript_output_hidden_state(self):
        pass

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

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

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

    @unittest.skip(reason="OneFormer does not use token embeddings")
    def test_resize_tokens_embeddings(self):
        pass

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

    def test_forward_signature(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            signature = inspect.signature(model.forward)
            
            arg_names = [*signature.parameters.keys()]

            expected_arg_names = ["pixel_values", "task_inputs"]
            self.assertListEqual(arg_names[:2], expected_arg_names)

    @slow
    def test_model_from_pretrained(self):
        for model_name in ["shi-labs/oneformer_ade20k_swin_tiny"]:
            model = OneFormerModel.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),
            "task_inputs": torch.randint(high=self.model_tester.vocab_size, size=(2, 77), device=torch_device).long(),
            "text_inputs": torch.randint(
                high=self.model_tester.vocab_size, size=(2, 6, 77), device=torch_device
            ).long(),
            "mask_labels": torch.randn((2, 150, *size), device=torch_device),
            "class_labels": torch.zeros(2, 150, device=torch_device).long(),
        }

        config = self.model_tester.get_config()
        config.is_training = True

        model = OneFormerForUniversalSegmentation(config).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_oneformer_model(config, **inputs, output_hidden_states=True)

    def test_attention_outputs(self):
        config, inputs = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config).to(torch_device)
            outputs = model(**inputs, output_attentions=True)
            self.assertTrue(outputs.attentions is not None)

    def test_initialization(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.contrastive_temperature = 1

        configs_no_init = _config_zero_init(config)
        for model_class in self.all_model_classes:
            model = model_class(config=configs_no_init)
            for name, param in model.named_parameters():
                if param.requires_grad:
                    self.assertIn(
                        ((param.data.mean() * 1e9).round() / 1e9).item(),
                        [0.0, 1.0],
                        msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                    )

    def test_training(self):
        if not self.model_tester.is_training:
            return
        
        model_class = self.all_model_classes[1]
        (
            config,
            pixel_values,
            task_inputs,
            text_inputs,
            pixel_mask,
            mask_labels,
            class_labels,
        ) = self.model_tester.prepare_config_and_inputs()
        config.is_training = True

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

        loss = model(
            pixel_values, task_inputs, text_inputs=text_inputs, mask_labels=mask_labels, class_labels=class_labels
        ).loss
        loss.backward()

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

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

        outputs = model(
            pixel_values, task_inputs, text_inputs=text_inputs, 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_class_predictions = outputs.transformer_decoder_class_predictions
        transformer_decoder_class_predictions.retain_grad()

        transformer_decoder_mask_predictions = outputs.transformer_decoder_mask_predictions
        transformer_decoder_mask_predictions.retain_grad()

        attentions = outputs.attentions[0][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_class_predictions.grad)
        self.assertIsNotNone(transformer_decoder_mask_predictions.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 OneFormerModelIntegrationTest(unittest.TestCase):
    @cached_property
    def model_checkpoints(self):
        return "shi-labs/oneformer_ade20k_swin_tiny"

    @cached_property
    def default_processor(self):
        return OneFormerProcessor.from_pretrained(self.model_checkpoints) if is_vision_available() else None

    def test_inference_no_head(self):
        model = OneFormerModel.from_pretrained(self.model_checkpoints).to(torch_device)
        processor = self.default_processor
        image = prepare_img()
        inputs = processor(image, ["semantic"], return_tensors="pt").to(torch_device)
        inputs_shape = inputs["pixel_values"].shape
        
        self.assertEqual(inputs_shape, (1, 3, 512, 682))

        task_inputs_shape = inputs["task_inputs"].shape
        
        self.assertEqual(task_inputs_shape, (1, 77))

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

        expected_slice_hidden_state = torch.tensor(
            [[0.2723, 0.8280, 0.6026], [1.2699, 1.1257, 1.1444], [1.1344, 0.6153, 0.4177]]
        ).to(torch_device)
        self.assertTrue(
            torch.allclose(
                outputs.encoder_hidden_states[-1][0, 0, :3, :3], expected_slice_hidden_state, atol=TOLERANCE
            )
        )

        expected_slice_hidden_state = torch.tensor(
            [[1.0581, 1.2276, 1.2003], [1.1903, 1.2925, 1.2862], [1.158, 1.2559, 1.3216]]
        ).to(torch_device)
        self.assertTrue(
            torch.allclose(
                outputs.pixel_decoder_hidden_states[0][0, 0, :3, :3], expected_slice_hidden_state, atol=TOLERANCE
            )
        )

        expected_slice_hidden_state = torch.tensor(
            [[3.0668, -1.1833, -5.1103], [3.344, -3.362, -5.1101], [2.6017, -4.3613, -4.1444]]
        ).to(torch_device)
        self.assertTrue(
            torch.allclose(
                outputs.transformer_decoder_class_predictions[0, :3, :3], expected_slice_hidden_state, atol=TOLERANCE
            )
        )

    def test_inference_universal_segmentation_head(self):
        model = OneFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints).to(torch_device).eval()
        processor = self.default_processor
        image = prepare_img()
        inputs = processor(image, ["semantic"], return_tensors="pt").to(torch_device)
        inputs_shape = inputs["pixel_values"].shape
        
        self.assertEqual(inputs_shape, (1, 3, 512, 682))

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

        
        masks_queries_logits = outputs.masks_queries_logits
        self.assertEqual(
            masks_queries_logits.shape,
            (1, model.config.num_queries, inputs_shape[-2] // 4, (inputs_shape[-1] + 2) // 4),
        )
        expected_slice = [[[3.1848, 4.2141, 4.1993], [2.9000, 3.5721, 3.6603], [2.5358, 3.0883, 3.6168]]]
        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.num_queries, model.config.num_labels + 1),
        )
        expected_slice = torch.tensor(
            [[3.0668, -1.1833, -5.1103], [3.3440, -3.3620, -5.1101], [2.6017, -4.3613, -4.1444]]
        ).to(torch_device)
        self.assertTrue(torch.allclose(class_queries_logits[0, :3, :3], expected_slice, atol=TOLERANCE))

    @require_torch_accelerator
    @require_torch_fp16
    def test_inference_fp16(self):
        model = (
            OneFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints)
            .to(torch_device, dtype=torch.float16)
            .eval()
        )
        processor = self.default_processor
        image = prepare_img()
        inputs = processor(image, ["semantic"], return_tensors="pt").to(torch_device, dtype=torch.float16)

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

    def test_with_segmentation_maps_and_loss(self):
        dummy_model = OneFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints)
        processor = self.default_processor
        processor.image_processor.num_text = dummy_model.config.num_queries - dummy_model.config.text_encoder_n_ctx
        dummy_model.config.is_training = True
        model = OneFormerForUniversalSegmentation(dummy_model.config).to(torch_device).eval()
        del dummy_model

        inputs = processor(
            [np.zeros((3, 512, 640)), np.zeros((3, 512, 640))],
            ["semantic", "semantic"],
            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["task_inputs"] = inputs["task_inputs"].to(torch_device)
        inputs["text_inputs"] = inputs["text_inputs"].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)