
'''simple docstring'''





import unittest

import numpy as np
import torch

from diffusers import KarrasVePipeline, KarrasVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()


class 						__lowerCAmelCase	(	unittest.TestCase      ):



      """simple docstring"""







      @property
      def 		snake_case_			(      self		:       Any				):
          torch.manual_seed(0				)
          __lowercase		:			int	    =       UNetaDModel(
              block_out_channels=(32, 64)				,       layers_per_block=2				,       sample_size=32				,       in_channels=3				,       out_channels=3				,       down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')				,       up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')				,       )
          return model




      def 		snake_case_			(      self		:       Optional[Any]				):
          __lowercase		:			Optional[Any]	    =       self.dummy_uncond_unet
          __lowercase		:			str	    =       KarrasVeScheduler()

          __lowercase		:			int	    =       KarrasVePipeline(unet=_snake_case				,       scheduler=_snake_case				)
          pipe.to(_snake_case				)
          pipe.set_progress_bar_config(disable=_snake_case				)

          __lowercase		:			Optional[int]	    =       torch.manual_seed(0				)
          __lowercase		:			int	    =       pipe(num_inference_steps=2				,       generator=_snake_case				,       output_type='''numpy'''				).images

          __lowercase		:			Optional[int]	    =       torch.manual_seed(0				)
          __lowercase		:			int	    =       pipe(num_inference_steps=2				,       generator=_snake_case				,       output_type='''numpy'''				,       return_dict=_snake_case				)[0]

          __lowercase		:			Optional[Any]	    =       image[0, -3:, -3:, -1]
          __lowercase		:			Dict	    =       image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 32, 32, 3)
          __lowercase		:			Tuple	    =       np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]				)

          assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2
          assert np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1E-2






@slow
@require_torch
class 						__lowerCAmelCase	(	unittest.TestCase      ):



      """simple docstring"""







      def 		snake_case_			(      self		:       str				):
          __lowercase		:			Dict	    =       '''google/ncsnpp-celebahq-256'''
          __lowercase		:			str	    =       UNetaDModel.from_pretrained(_snake_case				)
          __lowercase		:			List[Any]	    =       KarrasVeScheduler()

          __lowercase		:			List[Any]	    =       KarrasVePipeline(unet=_snake_case				,       scheduler=_snake_case				)
          pipe.to(_snake_case				)
          pipe.set_progress_bar_config(disable=_snake_case				)

          __lowercase		:			str	    =       torch.manual_seed(0				)
          __lowercase		:			Optional[int]	    =       pipe(num_inference_steps=20				,       generator=_snake_case				,       output_type='''numpy'''				).images

          __lowercase		:			Tuple	    =       image[0, -3:, -3:, -1]
          assert image.shape == (1, 256, 256, 3)
          __lowercase		:			Union[str, Any]	    =       np.array([0.5_78, 0.58_11, 0.59_24, 0.58_09, 0.5_87, 0.58_86, 0.58_61, 0.58_02, 0.5_86]				)

          assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2




import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, 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, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import UperNetForSemanticSegmentation
 from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor


class 						__lowerCAmelCase	:



      """simple docstring"""







      def __init__(      self		:       Dict				,       _snake_case		:       Optional[int]				,       _snake_case		:       List[Any]=13				,       _snake_case		:       int=32				,       _snake_case		:       int=3				,       _snake_case		:       Any=4				,       _snake_case		:       Optional[int]=[10, 20, 30, 40]				,       _snake_case		:       Optional[Any]=[2, 2, 3, 2]				,       _snake_case		:       Dict=True				,       _snake_case		:       List[Any]=True				,       _snake_case		:       int=37				,       _snake_case		:       Union[str, Any]="gelu"				,       _snake_case		:       Tuple=10				,       _snake_case		:       Tuple=0.02				,       _snake_case		:       List[str]=["stage2", "stage3", "stage4"]				,       _snake_case		:       Tuple=3				,       _snake_case		:       int=None				,       ):
          __lowercase		:			List[Any]	    =       parent
          __lowercase		:			Union[str, Any]	    =       batch_size
          __lowercase		:			Optional[Any]	    =       image_size
          __lowercase		:			Optional[Any]	    =       num_channels
          __lowercase		:			List[str]	    =       num_stages
          __lowercase		:			Union[str, Any]	    =       hidden_sizes
          __lowercase		:			Optional[Any]	    =       depths
          __lowercase		:			List[Any]	    =       is_training
          __lowercase		:			List[str]	    =       use_labels
          __lowercase		:			Tuple	    =       intermediate_size
          __lowercase		:			Union[str, Any]	    =       hidden_act
          __lowercase		:			Dict	    =       type_sequence_label_size
          __lowercase		:			Dict	    =       initializer_range
          __lowercase		:			str	    =       out_features
          __lowercase		:			Tuple	    =       num_labels
          __lowercase		:			Tuple	    =       scope
          __lowercase		:			Optional[Any]	    =       num_stages




      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			Optional[int]	    =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

          __lowercase		:			Dict	    =       None
          if self.use_labels:
              __lowercase		:			List[Any]	    =       ids_tensor([self.batch_size]				,       self.type_sequence_label_size				)

          __lowercase		:			Union[str, Any]	    =       self.get_config()

          return config, pixel_values, labels




      def 		snake_case_			(      self		:       List[str]				):
          return ConvNextConfig(
              num_channels=self.num_channels				,       num_stages=self.num_stages				,       hidden_sizes=self.hidden_sizes				,       depths=self.depths				,       is_training=self.is_training				,       intermediate_size=self.intermediate_size				,       hidden_act=self.hidden_act				,       out_features=self.out_features				,       )




      def 		snake_case_			(      self		:       int				):
          return UperNetConfig(
              backbone_config=self.get_backbone_config()				,       hidden_size=512				,       pool_scales=[1, 2, 3, 6]				,       use_auxiliary_head=_snake_case				,       auxiliary_loss_weight=0.4				,       auxiliary_in_channels=40				,       auxiliary_channels=256				,       auxiliary_num_convs=1				,       auxiliary_concat_input=_snake_case				,       loss_ignore_index=255				,       num_labels=self.num_labels				,       )




      def 		snake_case_			(      self		:       Tuple				,       _snake_case		:       Optional[Any]				,       _snake_case		:       List[Any]				,       _snake_case		:       List[str]				):
          __lowercase		:			Optional[int]	    =       UperNetForSemanticSegmentation(config=_snake_case				)
          model.to(_snake_case				)
          model.eval()
          __lowercase		:			Optional[Any]	    =       model(_snake_case				)
          self.parent.assertEqual(
              result.logits.shape				,       (self.batch_size, self.num_labels, self.image_size, self.image_size)				)




      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			Optional[int]	    =       self.prepare_config_and_inputs()
          (
              (
              __lowercase
          )    ,      (
              __lowercase
          )    ,      (
              __lowercase
          )    ,      
          )		:			Optional[int]	    =       config_and_inputs
          __lowercase		:			List[Any]	    =       {'''pixel_values''': pixel_values}
          return config, inputs_dict






@require_torch
class 						__lowerCAmelCase	(	lowerCAmelCase_ , lowerCAmelCase_ , unittest.TestCase      ):



      """simple docstring"""







      A__	:			List[str]        =    (UperNetForSemanticSegmentation,) if is_torch_available() else ()
      A__	:			Union[str, Any]        =    {'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
      A__	:			Union[str, Any]        =    False
      A__	:			Optional[Any]        =    False
      A__	:			int        =    False
      A__	:			Optional[int]        =    False
      A__	:			Optional[Any]        =    False
      A__	:			List[str]        =    False




      def 		snake_case_			(      self		:       Optional[Any]				):
          __lowercase		:			str	    =       UperNetModelTester(self				)
          __lowercase		:			List[Any]	    =       ConfigTester(self				,       config_class=_snake_case				,       has_text_modality=_snake_case				,       hidden_size=37				)




      def 		snake_case_			(      self		:       str				):
          self.create_and_test_config_common_properties()
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.create_and_test_config_with_num_labels()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()




      def 		snake_case_			(      self		:       str				):
          return




      def 		snake_case_			(      self		:       str				):
          __lowercase    ,      __lowercase		:			Optional[Any]	    =       self.model_tester.prepare_config_and_inputs_for_common()

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

              __lowercase		:			Dict	    =       ['''pixel_values''']
              self.assertListEqual(arg_names[:1]				,       _snake_case				)




      def 		snake_case_			(      self		:       List[str]				):
          __lowercase		:			int	    =       self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_semantic_segmentation(*_snake_case				)




      @unittest.skip(reason='''UperNet does not use inputs_embeds'''				)
      def 		snake_case_			(      self		:       List[Any]				):
          pass




      @unittest.skip(reason='''UperNet does not support input and output embeddings'''				)
      def 		snake_case_			(      self		:       List[Any]				):
          pass




      @unittest.skip(reason='''UperNet does not have a base model'''				)
      def 		snake_case_			(      self		:       str				):
          pass




      @unittest.skip(reason='''UperNet does not have a base model'''				)
      def 		snake_case_			(      self		:       int				):
          pass




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




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




      def 		snake_case_			(      self		:       List[Any]				):
          def check_hidden_states_output(_snake_case		:       int				,       _snake_case		:       Optional[int]				,       _snake_case		:       Dict				):
              __lowercase		:			Dict	    =       model_class(_snake_case				)
              model.to(_snake_case				)
              model.eval()

              with torch.no_grad():
                  __lowercase		:			Optional[int]	    =       model(**self._prepare_for_class(_snake_case				,       _snake_case				)				)

              __lowercase		:			Tuple	    =       outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

              __lowercase		:			Tuple	    =       self.model_tester.num_stages
              self.assertEqual(len(_snake_case				)				,       expected_num_stages + 1				)

              # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
              self.assertListEqual(
                  list(hidden_states[0].shape[-2:]				)				,       [self.model_tester.image_size // 4, self.model_tester.image_size // 4]				,       )

          __lowercase    ,      __lowercase		:			str	    =       self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              __lowercase		:			int	    =       True
              check_hidden_states_output(_snake_case				,       _snake_case				,       _snake_case				)

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              __lowercase		:			Any	    =       True

              check_hidden_states_output(_snake_case				,       _snake_case				,       _snake_case				)




      def 		snake_case_			(      self		:       Tuple				):
          __lowercase    ,      __lowercase		:			List[Any]	    =       self.model_tester.prepare_config_and_inputs_for_common()

          __lowercase		:			Dict	    =       _config_zero_init(_snake_case				)
          __lowercase		:			Dict	    =       _config_zero_init(configs_no_init.backbone_config				)
          for model_class in self.all_model_classes:
              __lowercase		:			Union[str, Any]	    =       model_class(config=_snake_case				)
              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'				,       )




      @unittest.skip(reason='''UperNet does not have tied weights'''				)
      def 		snake_case_			(      self		:       str				):
          pass




      @slow
      def 		snake_case_			(      self		:       Dict				):
          for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              __lowercase		:			List[Any]	    =       UperNetForSemanticSegmentation.from_pretrained(_snake_case				)
              self.assertIsNotNone(_snake_case				)






def     UpperCAmelCase_  (						) ->						Optional[int]:
    __lowercase		:			Optional[int]	    =       hf_hub_download(
        repo_id='''hf-internal-testing/fixtures_ade20k'''	,					repo_type='''dataset'''	,					filename='''ADE_val_00000001.jpg'''	)
    __lowercase		:			int	    =       Image.open(__lowerCAmelCase	).convert('''RGB'''	)
    return image





@require_torch
@require_vision
@slow
class 						__lowerCAmelCase	(	unittest.TestCase      ):



      """simple docstring"""







      def 		snake_case_			(      self		:       Optional[Any]				):
          __lowercase		:			List[Any]	    =       AutoImageProcessor.from_pretrained('''openmmlab/upernet-swin-tiny'''				)
          __lowercase		:			Optional[Any]	    =       UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-swin-tiny'''				).to(_snake_case				)

          __lowercase		:			Tuple	    =       prepare_img()
          __lowercase		:			Optional[int]	    =       processor(images=_snake_case				,       return_tensors='''pt'''				).to(_snake_case				)

          with torch.no_grad():
              __lowercase		:			Dict	    =       model(**_snake_case				)

          __lowercase		:			Optional[Any]	    =       torch.Size((1, model.config.num_labels, 512, 512)				)
          self.assertEqual(outputs.logits.shape				,       _snake_case				)

          __lowercase		:			List[Any]	    =       torch.tensor(
              [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]				).to(_snake_case				)
          self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]				,       _snake_case				,       atol=1E-4				)				)




      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			List[str]	    =       AutoImageProcessor.from_pretrained('''openmmlab/upernet-convnext-tiny'''				)
          __lowercase		:			Dict	    =       UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-convnext-tiny'''				).to(_snake_case				)

          __lowercase		:			Any	    =       prepare_img()
          __lowercase		:			Union[str, Any]	    =       processor(images=_snake_case				,       return_tensors='''pt'''				).to(_snake_case				)

          with torch.no_grad():
              __lowercase		:			Tuple	    =       model(**_snake_case				)

          __lowercase		:			Tuple	    =       torch.Size((1, model.config.num_labels, 512, 512)				)
          self.assertEqual(outputs.logits.shape				,       _snake_case				)

          __lowercase		:			List[Any]	    =       torch.tensor(
              [[-8.81_10, -8.81_10, -8.65_21], [-8.81_10, -8.81_10, -8.65_21], [-8.77_46, -8.77_46, -8.61_30]]				).to(_snake_case				)
          self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]				,       _snake_case				,       atol=1E-4				)				)


