






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 				_UpperCamelCase      :


      def __init__(						self       ,    __UpperCamelCase       ,    __UpperCamelCase=1_3       ,    __UpperCamelCase=3_2       ,    __UpperCamelCase=3       ,    __UpperCamelCase=4       ,    __UpperCamelCase=[1_0, 2_0, 3_0, 4_0]       ,    __UpperCamelCase=[2, 2, 3, 2]       ,    __UpperCamelCase=True       ,    __UpperCamelCase=True       ,    __UpperCamelCase=3_7       ,    __UpperCamelCase="gelu"       ,    __UpperCamelCase=1_0       ,    __UpperCamelCase=0.0_2       ,    __UpperCamelCase=["stage2", "stage3", "stage4"]       ,    __UpperCamelCase=3       ,    __UpperCamelCase=None       ,    )->						Optional[int]:
           __lowerCAmelCase             =       parent
           __lowerCAmelCase             =       batch_size
           __lowerCAmelCase             =       image_size
           __lowerCAmelCase             =       num_channels
           __lowerCAmelCase             =       num_stages
           __lowerCAmelCase             =       hidden_sizes
           __lowerCAmelCase             =       depths
           __lowerCAmelCase             =       is_training
           __lowerCAmelCase             =       use_labels
           __lowerCAmelCase             =       intermediate_size
           __lowerCAmelCase             =       hidden_act
           __lowerCAmelCase             =       type_sequence_label_size
           __lowerCAmelCase             =       initializer_range
           __lowerCAmelCase             =       out_features
           __lowerCAmelCase             =       num_labels
           __lowerCAmelCase             =       scope
           __lowerCAmelCase             =       num_stages


      def 			__UpperCAmelCase		(						self						)->						Optional[int]:
           __lowerCAmelCase             =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

           __lowerCAmelCase             =       None
           if self.use_labels:
                __lowerCAmelCase             =       ids_tensor([self.batch_size]       ,    self.type_sequence_label_size						)

           __lowerCAmelCase             =       self.get_config()

           return config, pixel_values, labels


      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           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 			__UpperCAmelCase		(						self						)->						List[Any]:
           return UperNetConfig(
               backbone_config=self.get_backbone_config()       ,    hidden_size=5_1_2       ,    pool_scales=[1, 2, 3, 6]       ,    use_auxiliary_head=__lowercase       ,    auxiliary_loss_weight=0.4       ,    auxiliary_in_channels=4_0       ,    auxiliary_channels=2_5_6       ,    auxiliary_num_convs=1       ,    auxiliary_concat_input=__lowercase       ,    loss_ignore_index=2_5_5       ,    num_labels=self.num_labels       ,    )


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase						)->						int:
           __lowerCAmelCase             =       UperNetForSemanticSegmentation(config=__lowercase						)
           model.to(__lowercase						)
           model.eval()
           __lowerCAmelCase             =       model(__lowercase						)
           self.parent.assertEqual(
               result.logits.shape       ,    (self.batch_size, self.num_labels, self.image_size, self.image_size)						)


      def 			__UpperCAmelCase		(						self						)->						Optional[int]:
           __lowerCAmelCase             =       self.prepare_config_and_inputs()
           (
               __lowerCAmelCase
           )             =       config_and_inputs
           __lowerCAmelCase             =       {'pixel_values': pixel_values}
           return config, inputs_dict



@require_torch
class 				_UpperCamelCase      (snake_case_				,		snake_case_				,		unittest.TestCase      ):
      snake_case_	=				(UperNetForSemanticSegmentation,) if is_torch_available() else ()
      snake_case_	=				{'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
      snake_case_	=				False
      snake_case_	=				False
      snake_case_	=				False
      snake_case_	=				False
      snake_case_	=				False
      snake_case_	=				False


      def 			__UpperCAmelCase		(						self						)->						Dict:
           __lowerCAmelCase             =       UperNetModelTester(self						)
           __lowerCAmelCase             =       ConfigTester(self       ,    config_class=__lowercase       ,    has_text_modality=__lowercase       ,    hidden_size=3_7						)


      def 			__UpperCAmelCase		(						self						)->						Tuple:
           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 			__UpperCAmelCase		(						self						)->						Any:
           return


      def 			__UpperCAmelCase		(						self						)->						Tuple:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs_for_common()

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

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


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*__lowercase						)


      @unittest.skip(reason="UperNet does not use inputs_embeds"						)
      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           pass


      @unittest.skip(reason="UperNet does not support input and output embeddings"						)
      def 			__UpperCAmelCase		(						self						)->						Any:
           pass


      @unittest.skip(reason="UperNet does not have a base model"						)
      def 			__UpperCAmelCase		(						self						)->						int:
           pass


      @unittest.skip(reason="UperNet does not have a base model"						)
      def 			__UpperCAmelCase		(						self						)->						Dict:
           pass


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


      @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."						)
      def 			__UpperCAmelCase		(						self						)->						Dict:
           pass


      def 			__UpperCAmelCase		(						self						)->						Union[str, Any]:







           def check_hidden_states_output(__UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase						):
                __lowerCAmelCase             =       model_class(__lowercase						)
                model.to(__lowercase						)
                model.eval()

                with torch.no_grad():
                     __lowerCAmelCase             =       model(**self._prepare_for_class(__lowercase       ,    __lowercase						)						)

                __lowerCAmelCase             =       outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                __lowerCAmelCase             =       self.model_tester.num_stages
                self.assertEqual(len(__lowercase						)       ,    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]       ,    )

           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs_for_common()

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

                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                __lowerCAmelCase             =       True

                check_hidden_states_output(__lowercase       ,    __lowercase       ,    __lowercase						)


      def 			__UpperCAmelCase		(						self						)->						Dict:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs_for_common()

           __lowerCAmelCase             =       _config_zero_init(__lowercase						)
           __lowerCAmelCase             =       _config_zero_init(configs_no_init.backbone_config						)
           for model_class in self.all_model_classes:
                __lowerCAmelCase             =       model_class(config=__lowercase						)
                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 			__UpperCAmelCase		(						self						)->						Dict:
           pass


      @slow
      def 			__UpperCAmelCase		(						self						)->						List[Any]:







           for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                __lowerCAmelCase             =       UperNetForSemanticSegmentation.from_pretrained(__lowercase						)
                self.assertIsNotNone(__lowercase						)





def    __lowerCAmelCase   (    ):
     __lowerCAmelCase             =       hf_hub_download(
         repo_id="hf-internal-testing/fixtures_ade20k"       ,			repo_type="dataset"       ,			filename="ADE_val_00000001.jpg"							)
     __lowerCAmelCase             =       Image.open(UpperCamelCase__							).convert("RGB"							)
     return image



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


      def 			__UpperCAmelCase		(						self						)->						Dict:
           __lowerCAmelCase             =       AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny"						)
           __lowerCAmelCase             =       UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny"						).to(__lowercase						)

           __lowerCAmelCase             =       prepare_img()
           __lowerCAmelCase             =       processor(images=__lowercase       ,    return_tensors="pt"						).to(__lowercase						)

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

           __lowerCAmelCase             =       torch.Size((1, model.config.num_labels, 5_1_2, 5_1_2)						)
           self.assertEqual(outputs.logits.shape       ,    __lowercase						)

           __lowerCAmelCase             =       torch.tensor(
               [[-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.4_7_9_7, -7.4_7_9_7, -7.3_0_6_8]]						).to(__lowercase						)
           self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]       ,    __lowercase       ,    atol=1e-4						)						)


      def 			__UpperCAmelCase		(						self						)->						Optional[int]:
           __lowerCAmelCase             =       AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny"						)
           __lowerCAmelCase             =       UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny"						).to(__lowercase						)

           __lowerCAmelCase             =       prepare_img()
           __lowerCAmelCase             =       processor(images=__lowercase       ,    return_tensors="pt"						).to(__lowercase						)

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

           __lowerCAmelCase             =       torch.Size((1, model.config.num_labels, 5_1_2, 5_1_2)						)
           self.assertEqual(outputs.logits.shape       ,    __lowercase						)

           __lowerCAmelCase             =       torch.tensor(
               [[-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.7_7_4_6, -8.7_7_4_6, -8.6_1_3_0]]						).to(__lowercase						)
           self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]       ,    __lowercase       ,    atol=1e-4						)						)










'''simple docstring'''
import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp


def       _lowercase							(					):
 __A	:		List[Any]	  =      ArgumentParser(
     description=(
         'PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes'
     )      )

 # Optional arguments for the launch helper
 parser.add_argument('--num_cores',					type=UpperCamelCase__,					default=1,					help='Number of TPU cores to use (1 or 8).'      )

 # positional
 parser.add_argument(
     'training_script',					type=UpperCamelCase__,					help=(
         'The full path to the single TPU training '
         'program/script to be launched in parallel, '
         'followed by all the arguments for the '
         'training script'
     ),					)

 # rest from the training program
 parser.add_argument('training_script_args',					nargs=UpperCamelCase__      )

 return parser.parse_args()


def       _lowercase							(					):
 __A	:		str	  =      parse_args()

 # Import training_script as a module.
 __A	:		Tuple	  =      Path(args.training_script      )
 sys.path.append(str(script_fpath.parent.resolve()      )      )
 __A	:		str	  =      script_fpath.stem
 __A	:		Dict	  =      importlib.import_module(UpperCamelCase__      )

 # Patch sys.argv
 __A	:		Optional[Any]	  =      [args.training_script] + args.training_script_args + ['--tpu_num_cores', str(args.num_cores      )]

 xmp.spawn(mod._mp_fn,					args=(),					nprocs=args.num_cores      )


if __name__ == "__main__":
 main()






