init

transformers/tests/models/seggpt/test_image_processing_seggpt.py

@@ -0,0 +1,310 @@
+# Copyright 2024 HuggingFace Inc.
+#
+# 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.
+
+import unittest
+
+import numpy as np
+from datasets import load_dataset
+
+from transformers.testing_utils import require_torch, require_vision, slow
+from transformers.utils import is_torch_available, is_vision_available
+
+from ...test_image_processing_common import ImageProcessingTestMixin, prepare_image_inputs
+
+
+if is_torch_available():
+    import torch
+
+    from transformers.models.seggpt.modeling_seggpt import SegGptImageSegmentationOutput
+
+if is_vision_available():
+    from PIL import Image
+
+    from transformers import SegGptImageProcessor
+
+
+class SegGptImageProcessingTester:
+    def __init__(
+        self,
+        parent,
+        batch_size=7,
+        num_channels=3,
+        image_size=18,
+        min_resolution=30,
+        max_resolution=400,
+        do_resize=True,
+        size=None,
+        do_normalize=True,
+        image_mean=[0.5, 0.5, 0.5],
+        image_std=[0.5, 0.5, 0.5],
+    ):
+        size = size if size is not None else {"height": 18, "width": 18}
+        self.parent = parent
+        self.batch_size = batch_size
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.min_resolution = min_resolution
+        self.max_resolution = max_resolution
+        self.do_resize = do_resize
+        self.size = size
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean
+        self.image_std = image_std
+
+    def prepare_image_processor_dict(self):
+        return {
+            "image_mean": self.image_mean,
+            "image_std": self.image_std,
+            "do_normalize": self.do_normalize,
+            "do_resize": self.do_resize,
+            "size": self.size,
+        }
+
+    def expected_output_image_shape(self, images):
+        return self.num_channels, self.size["height"], self.size["width"]
+
+    def expected_post_processed_shape(self):
+        return self.size["height"] // 2, self.size["width"]
+
+    def get_fake_image_segmentation_output(self):
+        torch.manual_seed(42)
+        return SegGptImageSegmentationOutput(
+            pred_masks=torch.rand(self.batch_size, self.num_channels, self.size["height"], self.size["width"])
+        )
+
+    def prepare_image_inputs(self, equal_resolution=False, numpify=False, torchify=False):
+        return prepare_image_inputs(
+            batch_size=self.batch_size,
+            num_channels=self.num_channels,
+            min_resolution=self.min_resolution,
+            max_resolution=self.max_resolution,
+            equal_resolution=equal_resolution,
+            numpify=numpify,
+            torchify=torchify,
+        )
+
+
+def prepare_mask():
+    ds = load_dataset("EduardoPacheco/seggpt-example-data")["train"]
+    return ds[0]["mask"].convert("L")
+
+
+def prepare_img():
+    ds = load_dataset("EduardoPacheco/seggpt-example-data")["train"]
+    images = [image.convert("RGB") for image in ds["image"]]
+    masks = [image.convert("RGB") for image in ds["mask"]]
+    return images, masks
+
+
+@require_torch
+@require_vision
+class SegGptImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase):
+    image_processing_class = SegGptImageProcessor if is_vision_available() else None
+
+    def setUp(self):
+        super().setUp()
+        self.image_processor_tester = SegGptImageProcessingTester(self)
+
+    @property
+    def image_processor_dict(self):
+        return self.image_processor_tester.prepare_image_processor_dict()
+
+    def test_image_processor_properties(self):
+        image_processing = self.image_processing_class(**self.image_processor_dict)
+        self.assertTrue(hasattr(image_processing, "image_mean"))
+        self.assertTrue(hasattr(image_processing, "image_std"))
+        self.assertTrue(hasattr(image_processing, "do_normalize"))
+        self.assertTrue(hasattr(image_processing, "do_resize"))
+        self.assertTrue(hasattr(image_processing, "size"))
+
+    def test_image_processor_from_dict_with_kwargs(self):
+        image_processor = self.image_processing_class.from_dict(self.image_processor_dict)
+        self.assertEqual(image_processor.size, {"height": 18, "width": 18})
+
+        image_processor = self.image_processing_class.from_dict(self.image_processor_dict, size=42)
+        self.assertEqual(image_processor.size, {"height": 42, "width": 42})
+
+    def test_image_processor_palette(self):
+        num_labels = 3
+        image_processing = self.image_processing_class(**self.image_processor_dict)
+        palette = image_processing.get_palette(num_labels)
+        self.assertEqual(len(palette), num_labels + 1)
+        self.assertEqual(palette[0], (0, 0, 0))
+
+    def test_mask_equivalence(self):
+        image_processor = SegGptImageProcessor()
+
+        mask_binary = prepare_mask()
+        mask_rgb = mask_binary.convert("RGB")
+
+        inputs_binary = image_processor(images=None, prompt_masks=mask_binary, return_tensors="pt")
+        inputs_rgb = image_processor(images=None, prompt_masks=mask_rgb, return_tensors="pt", do_convert_rgb=False)
+
+        self.assertTrue((inputs_binary["prompt_masks"] == inputs_rgb["prompt_masks"]).all().item())
+
+    def test_mask_to_rgb(self):
+        image_processing = self.image_processing_class(**self.image_processor_dict)
+        mask = prepare_mask()
+        mask = np.array(mask)
+        mask = (mask > 0).astype(np.uint8)
+
+        def check_two_colors(image, color1=(0, 0, 0), color2=(255, 255, 255)):
+            pixels = image.transpose(1, 2, 0).reshape(-1, 3)
+            unique_colors = np.unique(pixels, axis=0)
+            if len(unique_colors) == 2 and (color1 in unique_colors) and (color2 in unique_colors):
+                return True
+            else:
+                return False
+
+        num_labels = 1
+        palette = image_processing.get_palette(num_labels)
+
+        # Should only duplicate repeat class indices map, hence only (0,0,0) and (1,1,1)
+        mask_duplicated = image_processing.mask_to_rgb(mask)
+        # Mask using palette, since only 1 class is present we have colors (0,0,0) and (255,255,255)
+        mask_painted = image_processing.mask_to_rgb(mask, palette=palette)
+
+        self.assertTrue(check_two_colors(mask_duplicated, color2=(1, 1, 1)))
+        self.assertTrue(check_two_colors(mask_painted, color2=(255, 255, 255)))
+
+    def test_post_processing_semantic_segmentation(self):
+        image_processor = self.image_processing_class(**self.image_processor_dict)
+        outputs = self.image_processor_tester.get_fake_image_segmentation_output()
+        post_processed = image_processor.post_process_semantic_segmentation(outputs)
+
+        self.assertEqual(len(post_processed), self.image_processor_tester.batch_size)
+
+        expected_semantic_map_shape = self.image_processor_tester.expected_post_processed_shape()
+        self.assertEqual(post_processed[0].shape, expected_semantic_map_shape)
+
+    @slow
+    def test_pixel_values(self):
+        images, masks = prepare_img()
+        input_image = images[1]
+        prompt_image = images[0]
+        prompt_mask = masks[0]
+
+        image_processor = SegGptImageProcessor.from_pretrained("BAAI/seggpt-vit-large")
+
+        inputs = image_processor(
+            images=input_image,
+            prompt_images=prompt_image,
+            prompt_masks=prompt_mask,
+            return_tensors="pt",
+            do_convert_rgb=False,
+        )
+
+        # Verify pixel values
+        expected_prompt_pixel_values = torch.tensor(
+            [
+                [[-0.6965, -0.6965, -0.6965], [-0.6965, -0.6965, -0.6965], [-0.6965, -0.6965, -0.6965]],
+                [[1.6583, 1.6583, 1.6583], [1.6583, 1.6583, 1.6583], [1.6583, 1.6583, 1.6583]],
+                [[2.3088, 2.3088, 2.3088], [2.3088, 2.3088, 2.3088], [2.3088, 2.3088, 2.3088]],
+            ]
+        )
+
+        expected_pixel_values = torch.tensor(
+            [
+                [[1.6324, 1.6153, 1.5810], [1.6153, 1.5982, 1.5810], [1.5810, 1.5639, 1.5639]],
+                [[1.2731, 1.2556, 1.2206], [1.2556, 1.2381, 1.2031], [1.2206, 1.2031, 1.1681]],
+                [[1.6465, 1.6465, 1.6465], [1.6465, 1.6465, 1.6465], [1.6291, 1.6291, 1.6291]],
+            ]
+        )
+
+        expected_prompt_masks = torch.tensor(
+            [
+                [[-2.1179, -2.1179, -2.1179], [-2.1179, -2.1179, -2.1179], [-2.1179, -2.1179, -2.1179]],
+                [[-2.0357, -2.0357, -2.0357], [-2.0357, -2.0357, -2.0357], [-2.0357, -2.0357, -2.0357]],
+                [[-1.8044, -1.8044, -1.8044], [-1.8044, -1.8044, -1.8044], [-1.8044, -1.8044, -1.8044]],
+            ]
+        )
+
+        torch.testing.assert_close(inputs.pixel_values[0, :, :3, :3], expected_pixel_values, rtol=1e-4, atol=1e-4)
+        torch.testing.assert_close(
+            inputs.prompt_pixel_values[0, :, :3, :3], expected_prompt_pixel_values, rtol=1e-4, atol=1e-4
+        )
+        torch.testing.assert_close(inputs.prompt_masks[0, :, :3, :3], expected_prompt_masks, rtol=1e-4, atol=1e-4)
+
+    def test_prompt_mask_equivalence(self):
+        image_processor = self.image_processing_class(**self.image_processor_dict)
+        image_size = self.image_processor_tester.image_size
+
+        # Single Mask Examples
+        expected_single_shape = [1, 3, image_size, image_size]
+
+        # Single Semantic Map (2D)
+        image_np_2d = np.ones((image_size, image_size))
+        image_pt_2d = torch.ones((image_size, image_size))
+        image_pil_2d = Image.fromarray(image_np_2d)
+
+        inputs_np_2d = image_processor(images=None, prompt_masks=image_np_2d, return_tensors="pt")
+        inputs_pt_2d = image_processor(images=None, prompt_masks=image_pt_2d, return_tensors="pt")
+        inputs_pil_2d = image_processor(images=None, prompt_masks=image_pil_2d, return_tensors="pt")
+
+        self.assertTrue((inputs_np_2d["prompt_masks"] == inputs_pt_2d["prompt_masks"]).all().item())
+        self.assertTrue((inputs_np_2d["prompt_masks"] == inputs_pil_2d["prompt_masks"]).all().item())
+        self.assertEqual(list(inputs_np_2d["prompt_masks"].shape), expected_single_shape)
+
+        # Single RGB Images (3D)
+        image_np_3d = np.ones((3, image_size, image_size))
+        image_pt_3d = torch.ones((3, image_size, image_size))
+        image_pil_3d = Image.fromarray(image_np_3d.transpose(1, 2, 0).astype(np.uint8))
+
+        inputs_np_3d = image_processor(
+            images=None, prompt_masks=image_np_3d, return_tensors="pt", do_convert_rgb=False
+        )
+        inputs_pt_3d = image_processor(
+            images=None, prompt_masks=image_pt_3d, return_tensors="pt", do_convert_rgb=False
+        )
+        inputs_pil_3d = image_processor(
+            images=None, prompt_masks=image_pil_3d, return_tensors="pt", do_convert_rgb=False
+        )
+
+        self.assertTrue((inputs_np_3d["prompt_masks"] == inputs_pt_3d["prompt_masks"]).all().item())
+        self.assertTrue((inputs_np_3d["prompt_masks"] == inputs_pil_3d["prompt_masks"]).all().item())
+        self.assertEqual(list(inputs_np_3d["prompt_masks"].shape), expected_single_shape)
+
+        # Batched Examples
+        expected_batched_shape = [2, 3, image_size, image_size]
+
+        # Batched Semantic Maps (3D)
+        image_np_2d_batched = np.ones((2, image_size, image_size))
+        image_pt_2d_batched = torch.ones((2, image_size, image_size))
+
+        inputs_np_2d_batched = image_processor(images=None, prompt_masks=image_np_2d_batched, return_tensors="pt")
+        inputs_pt_2d_batched = image_processor(images=None, prompt_masks=image_pt_2d_batched, return_tensors="pt")
+
+        self.assertTrue((inputs_np_2d_batched["prompt_masks"] == inputs_pt_2d_batched["prompt_masks"]).all().item())
+        self.assertEqual(list(inputs_np_2d_batched["prompt_masks"].shape), expected_batched_shape)
+
+        # Batched RGB images
+        image_np_4d = np.ones((2, 3, image_size, image_size))
+        image_pt_4d = torch.ones((2, 3, image_size, image_size))
+
+        inputs_np_4d = image_processor(
+            images=None, prompt_masks=image_np_4d, return_tensors="pt", do_convert_rgb=False
+        )
+        inputs_pt_4d = image_processor(
+            images=None, prompt_masks=image_pt_4d, return_tensors="pt", do_convert_rgb=False
+        )
+
+        self.assertTrue((inputs_np_4d["prompt_masks"] == inputs_pt_4d["prompt_masks"]).all().item())
+        self.assertEqual(list(inputs_np_4d["prompt_masks"].shape), expected_batched_shape)
+
+        # Comparing Single and Batched Examples
+        self.assertTrue((inputs_np_2d["prompt_masks"][0] == inputs_np_3d["prompt_masks"][0]).all().item())
+        self.assertTrue((inputs_np_2d_batched["prompt_masks"][0] == inputs_np_2d["prompt_masks"][0]).all().item())
+        self.assertTrue((inputs_np_2d_batched["prompt_masks"][0] == inputs_np_3d["prompt_masks"][0]).all().item())
+        self.assertTrue((inputs_np_2d_batched["prompt_masks"][0] == inputs_np_4d["prompt_masks"][0]).all().item())
+        self.assertTrue((inputs_np_2d_batched["prompt_masks"][0] == inputs_np_3d["prompt_masks"][0]).all().item())

transformers/tests/models/seggpt/test_modeling_seggpt.py

@@ -0,0 +1,475 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# 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 SegGpt model."""
+
+import inspect
+import math
+import unittest
+from functools import cached_property
+
+from datasets import load_dataset
+
+from transformers import SegGptConfig
+from transformers.testing_utils import (
+    Expectations,
+    require_torch,
+    require_vision,
+    slow,
+    torch_device,
+)
+from transformers.utils import is_torch_available, is_vision_available
+
+from ...test_configuration_common import ConfigTester
+from ...test_modeling_common import ModelTesterMixin, floats_tensor
+from ...test_pipeline_mixin import PipelineTesterMixin
+
+
+if is_torch_available():
+    import torch
+    from torch import nn
+
+    from transformers import SegGptForImageSegmentation, SegGptModel
+    from transformers.models.seggpt.modeling_seggpt import SegGptLoss
+
+
+if is_vision_available():
+    from transformers import SegGptImageProcessor
+
+
+class SegGptModelTester:
+    def __init__(
+        self,
+        parent,
+        batch_size=2,
+        image_size=30,
+        patch_size=2,
+        num_channels=3,
+        is_training=False,
+        use_labels=True,
+        hidden_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        initializer_range=0.02,
+        mlp_ratio=2.0,
+        merge_index=0,
+        intermediate_hidden_state_indices=[1],
+        pretrain_image_size=10,
+        decoder_hidden_size=10,
+    ):
+        self.parent = parent
+        self.batch_size = batch_size
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.is_training = is_training
+        self.use_labels = use_labels
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.initializer_range = initializer_range
+        self.mlp_ratio = mlp_ratio
+        self.merge_index = merge_index
+        self.intermediate_hidden_state_indices = intermediate_hidden_state_indices
+        self.pretrain_image_size = pretrain_image_size
+        self.decoder_hidden_size = decoder_hidden_size
+
+        # in SegGpt, the seq length equals the number of patches (we don't use the [CLS] token)
+        num_patches = (image_size // patch_size) ** 2
+        self.seq_length = num_patches
+
+    def prepare_config_and_inputs(self):
+        pixel_values = floats_tensor([self.batch_size, self.num_channels, self.image_size // 2, self.image_size])
+        prompt_pixel_values = floats_tensor(
+            [self.batch_size, self.num_channels, self.image_size // 2, self.image_size]
+        )
+        prompt_masks = floats_tensor([self.batch_size, self.num_channels, self.image_size // 2, self.image_size])
+
+        labels = None
+        if self.use_labels:
+            labels = floats_tensor([self.batch_size, self.num_channels, self.image_size // 2, self.image_size])
+
+        config = self.get_config()
+
+        return config, pixel_values, prompt_pixel_values, prompt_masks, labels
+
+    def get_config(self):
+        return SegGptConfig(
+            image_size=self.image_size,
+            patch_size=self.patch_size,
+            num_channels=self.num_channels,
+            hidden_size=self.hidden_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            hidden_act=self.hidden_act,
+            hidden_dropout_prob=self.hidden_dropout_prob,
+            initializer_range=self.initializer_range,
+            mlp_ratio=self.mlp_ratio,
+            merge_index=self.merge_index,
+            intermediate_hidden_state_indices=self.intermediate_hidden_state_indices,
+            pretrain_image_size=self.pretrain_image_size,
+            decoder_hidden_size=self.decoder_hidden_size,
+        )
+
+    def create_and_check_model(self, config, pixel_values, prompt_pixel_values, prompt_masks, labels):
+        model = SegGptModel(config=config)
+        model.to(torch_device)
+        model.eval()
+        result = model(pixel_values, prompt_pixel_values, prompt_masks)
+        self.parent.assertEqual(
+            result.last_hidden_state.shape,
+            (
+                self.batch_size,
+                self.image_size // self.patch_size,
+                self.image_size // self.patch_size,
+                self.hidden_size,
+            ),
+        )
+
+    def prepare_config_and_inputs_for_common(self):
+        config_and_inputs = self.prepare_config_and_inputs()
+        (
+            config,
+            pixel_values,
+            prompt_pixel_values,
+            prompt_masks,
+            labels,
+        ) = config_and_inputs
+        inputs_dict = {
+            "pixel_values": pixel_values,
+            "prompt_pixel_values": prompt_pixel_values,
+            "prompt_masks": prompt_masks,
+        }
+        return config, inputs_dict
+
+
+@require_torch
+class SegGptModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
+    """
+    Here we also overwrite some of the tests of test_modeling_common.py, as SegGpt does not use input_ids, inputs_embeds,
+    attention_mask and seq_length.
+    """
+
+    all_model_classes = (SegGptModel, SegGptForImageSegmentation) if is_torch_available() else ()
+    fx_compatible = False
+
+    test_pruning = False
+    test_resize_embeddings = False
+    test_head_masking = False
+    test_torchscript = False
+    test_torch_exportable = True
+
+    pipeline_model_mapping = (
+        {"feature-extraction": SegGptModel, "mask-generation": SegGptModel} if is_torch_available() else {}
+    )
+
+    def setUp(self):
+        self.model_tester = SegGptModelTester(self)
+        self.config_tester = ConfigTester(self, config_class=SegGptConfig, has_text_modality=False)
+
+    def test_config(self):
+        self.config_tester.run_common_tests()
+
+    @unittest.skip(reason="SegGpt does not use inputs_embeds")
+    def test_inputs_embeds(self):
+        pass
+
+    def test_model_get_set_embeddings(self):
+        config, _ = self.model_tester.prepare_config_and_inputs_for_common()
+
+        for model_class in self.all_model_classes:
+            model = model_class(config)
+            self.assertIsInstance(model.get_input_embeddings(), (nn.Module))
+
+    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)
+            # signature.parameters is an OrderedDict => so arg_names order is deterministic
+            arg_names = [*signature.parameters.keys()]
+
+            expected_arg_names = ["pixel_values", "prompt_pixel_values", "prompt_masks"]
+            self.assertListEqual(arg_names[:3], expected_arg_names)
+
+    def test_model(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_model(*config_and_inputs)
+
+    def test_hidden_states_output(self):
+        def check_hidden_states_output(inputs_dict, config, model_class):
+            model = model_class(config)
+            model.to(torch_device)
+            model.eval()
+
+            with torch.no_grad():
+                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
+
+            hidden_states = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states
+
+            expected_num_layers = getattr(
+                self.model_tester, "expected_num_hidden_layers", self.model_tester.num_hidden_layers + 1
+            )
+            self.assertEqual(len(hidden_states), expected_num_layers)
+
+            patch_height = patch_width = config.image_size // config.patch_size
+
+            self.assertListEqual(
+                list(hidden_states[0].shape[-3:]),
+                [patch_height, patch_width, self.model_tester.hidden_size],
+            )
+
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+
+        for model_class in self.all_model_classes:
+            inputs_dict["output_hidden_states"] = True
+            check_hidden_states_output(inputs_dict, config, model_class)
+
+            # check that output_hidden_states also work using config
+            del inputs_dict["output_hidden_states"]
+            config.output_hidden_states = True
+
+            check_hidden_states_output(inputs_dict, config, model_class)
+
+    def test_batching_equivalence(self):
+        def recursive_check(batched_object, single_row_object, model_name, key):
+            if isinstance(batched_object, (list, tuple)):
+                for batched_object_value, single_row_object_value in zip(batched_object, single_row_object):
+                    recursive_check(batched_object_value, single_row_object_value, model_name, key)
+            else:
+                batched_row = batched_object[:1]
+                self.assertFalse(
+                    torch.isnan(batched_row).any(), f"Batched output has `nan` in {model_name} for key={key}"
+                )
+                self.assertFalse(
+                    torch.isinf(batched_row).any(), f"Batched output has `inf` in {model_name} for key={key}"
+                )
+                self.assertFalse(
+                    torch.isnan(single_row_object).any(), f"Single row output has `nan` in {model_name} for key={key}"
+                )
+                self.assertFalse(
+                    torch.isinf(single_row_object).any(), f"Single row output has `inf` in {model_name} for key={key}"
+                )
+                self.assertTrue(
+                    torch.max(torch.abs(batched_row - single_row_object)) <= 1e-03,
+                    msg=(
+                        f"Batched and Single row outputs are not equal in {model_name} for key={key}. "
+                        f"Difference={torch.max(torch.abs(batched_row - single_row_object))}."
+                    ),
+                )
+
+        config, batched_input = self.model_tester.prepare_config_and_inputs_for_common()
+
+        for model_class in self.all_model_classes:
+            config.output_hidden_states = True
+
+            model_name = model_class.__name__
+            batched_input_prepared = self._prepare_for_class(batched_input, model_class)
+            model = model_class(config).to(torch_device).eval()
+
+            batch_size = self.model_tester.batch_size
+            single_row_input = {}
+            for key, value in batched_input_prepared.items():
+                if isinstance(value, torch.Tensor) and value.shape[0] % batch_size == 0:
+                    single_batch_shape = value.shape[0] // batch_size
+                    single_row_input[key] = value[:single_batch_shape]
+
+            with torch.no_grad():
+                model_batched_output = model(**batched_input_prepared)
+                model_row_output = model(**single_row_input)
+
+            for key in model_batched_output:
+                # the first hidden state in SegGPT has weird hack of adding first half of batch with second half
+                if key == "hidden_states":
+                    model_batched_output[key] = model_batched_output[key][1:]
+                    model_row_output[key] = model_row_output[key][1:]
+                recursive_check(model_batched_output[key], model_row_output[key], model_name, key)
+
+    def test_seggpt_loss(self):
+        torch.manual_seed(100)
+        config = self.model_tester.get_config()
+
+        prompt_masks = torch.rand(1, config.num_channels, config.image_size, config.image_size)
+        label = torch.rand(1, config.num_channels, config.image_size, config.image_size)
+        pred_masks = torch.rand(1, config.num_channels, config.image_size * 2, config.image_size)
+        # seq_len x 2 because the loss concatenates prompt_masks and labels as pred_masks is concatenated
+        bool_masked_pos = torch.rand(1, self.model_tester.seq_length * 2) > 0.5
+
+        loss = SegGptLoss(config)
+        loss_value = loss(prompt_masks, pred_masks, label, bool_masked_pos)
+        expected_loss_value = torch.tensor(0.3340)
+
+        torch.testing.assert_close(loss_value, expected_loss_value, rtol=1e-4, atol=1e-4)
+
+    @slow
+    def test_model_from_pretrained(self):
+        model_name = "BAAI/seggpt-vit-large"
+        model = SegGptModel.from_pretrained(model_name)
+        self.assertIsNotNone(model)
+
+
+def prepare_img():
+    ds = load_dataset("EduardoPacheco/seggpt-example-data")["train"]
+    images = [image.convert("RGB") for image in ds["image"]]
+    masks = [image.convert("RGB") for image in ds["mask"]]
+    return images, masks
+
+
+def prepare_bool_masked_pos(config: SegGptConfig):
+    num_patches = math.prod([i // config.patch_size for i in config.image_size])
+    mask_ratio = 0.75
+    torch.manual_seed(2)
+    num_masked_patches = int(num_patches * mask_ratio)
+    shuffle_idx = torch.randperm(num_patches)
+    bool_masked_pos = torch.FloatTensor([0] * (num_patches - num_masked_patches) + [1] * num_masked_patches)[
+        shuffle_idx
+    ]
+    bool_masked_pos = bool_masked_pos.unsqueeze(0).bool()
+
+    return bool_masked_pos
+
+
+@require_torch
+@require_vision
+class SegGptModelIntegrationTest(unittest.TestCase):
+    @cached_property
+    def default_image_processor(self):
+        return SegGptImageProcessor.from_pretrained("BAAI/seggpt-vit-large") if is_vision_available() else None
+
+    @slow
+    def test_one_shot_inference(self):
+        model = SegGptForImageSegmentation.from_pretrained("BAAI/seggpt-vit-large").to(torch_device)
+
+        image_processor = self.default_image_processor
+
+        images, masks = prepare_img()
+        input_image = images[1]
+        prompt_image = images[0]
+        prompt_mask = masks[0]
+
+        inputs = image_processor(
+            images=input_image,
+            prompt_images=prompt_image,
+            prompt_masks=prompt_mask,
+            return_tensors="pt",
+            do_convert_rgb=False,
+        )
+
+        inputs = inputs.to(torch_device)
+        # forward pass
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        # verify the logits
+        expected_shape = torch.Size((1, 3, 896, 448))
+        self.assertEqual(outputs.pred_masks.shape, expected_shape)
+
+        expectations = Expectations(
+            {
+                (None, None): [
+                    [[-2.1208, -2.1190, -2.1198], [-2.1237, -2.1228, -2.1227], [-2.1232, -2.1226, -2.1228]],
+                    [[-2.0405, -2.0396, -2.0403], [-2.0434, -2.0434, -2.0433], [-2.0428, -2.0432, -2.0434]],
+                    [[-1.8102, -1.8088, -1.8099], [-1.8131, -1.8126, -1.8129], [-1.8130, -1.8128, -1.8131]],
+                ],
+                ("cuda", 8): [
+                    [[-2.1208, -2.1189, -2.1198], [-2.1236, -2.1229, -2.1230], [-2.1233, -2.1227, -2.1228]],
+                    [[-2.0408, -2.0398, -2.0405], [-2.0435, -2.0437, -2.0438], [-2.0431, -2.0435, -2.0436]],
+                    [[-1.8101, -1.8086, -1.8098], [-1.8129, -1.8126, -1.8130], [-1.8128, -1.8128, -1.8130]],
+                ],
+            }
+        )
+        expected_slice = torch.tensor(expectations.get_expectation()).to(torch_device)
+
+        torch.testing.assert_close(outputs.pred_masks[0, :, :3, :3], expected_slice, rtol=2e-4, atol=2e-4)
+
+        result = image_processor.post_process_semantic_segmentation(outputs, [input_image.size[::-1]])[0]
+
+        result_expected_shape = torch.Size((170, 297))
+        expected_area = 1082
+        area = (result > 0).sum().item()
+        self.assertEqual(result.shape, result_expected_shape)
+        self.assertEqual(area, expected_area)
+
+    @slow
+    def test_few_shot_inference(self):
+        model = SegGptForImageSegmentation.from_pretrained("BAAI/seggpt-vit-large").to(torch_device)
+        image_processor = self.default_image_processor
+
+        images, masks = prepare_img()
+        input_images = [images[1]] * 2
+        prompt_images = [images[0], images[2]]
+        prompt_masks = [masks[0], masks[2]]
+
+        inputs = image_processor(
+            images=input_images,
+            prompt_images=prompt_images,
+            prompt_masks=prompt_masks,
+            return_tensors="pt",
+            do_convert_rgb=False,
+        )
+
+        inputs = {k: v.to(torch_device) for k, v in inputs.items()}
+        with torch.no_grad():
+            outputs = model(**inputs, feature_ensemble=True)
+
+        expected_shape = torch.Size((2, 3, 896, 448))
+        expected_slice = torch.tensor(
+            [
+                [[-2.1201, -2.1192, -2.1189], [-2.1217, -2.1210, -2.1204], [-2.1216, -2.1202, -2.1194]],
+                [[-2.0393, -2.0390, -2.0387], [-2.0402, -2.0402, -2.0397], [-2.0400, -2.0394, -2.0388]],
+                [[-1.8083, -1.8076, -1.8077], [-1.8105, -1.8102, -1.8099], [-1.8105, -1.8095, -1.8090]],
+            ]
+        ).to(torch_device)
+
+        self.assertEqual(outputs.pred_masks.shape, expected_shape)
+        torch.testing.assert_close(outputs.pred_masks[0, :, 448:451, :3], expected_slice, rtol=4e-4, atol=4e-4)
+
+    @slow
+    def test_one_shot_with_label(self):
+        model = SegGptForImageSegmentation.from_pretrained("BAAI/seggpt-vit-large").to(torch_device)
+
+        image_processor = self.default_image_processor
+
+        images, masks = prepare_img()
+
+        input_image = images[1]
+        label = masks[1]
+        prompt_image = images[0]
+        prompt_mask = masks[0]
+
+        inputs = image_processor(
+            images=input_image,
+            prompt_masks=prompt_mask,
+            prompt_images=prompt_image,
+            return_tensors="pt",
+            do_convert_rgb=False,
+        ).to(torch_device)
+
+        labels = image_processor(images=None, prompt_masks=label, return_tensors="pt", do_convert_rgb=False)[
+            "prompt_masks"
+        ].to(torch_device)
+
+        bool_masked_pos = prepare_bool_masked_pos(model.config).to(torch_device)
+
+        with torch.no_grad():
+            outputs = model(**inputs, labels=labels, bool_masked_pos=bool_masked_pos)
+
+        expected_loss = torch.tensor(0.0074).to(torch_device)
+        torch.testing.assert_close(outputs.loss, expected_loss, rtol=1e-4, atol=1e-4)