




"""simple docstring"""






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





						def __init__(							self      :	int     ,a__      :	Any     ,a__      :	Dict=13     ,a__      :	Optional[Any]=32     ,a__      :	List[str]=3     ,a__      :	int=4     ,a__      :	str=[10, 20, 30, 40]     ,a__      :	Union[str, Any]=[2, 2, 3, 2]     ,a__      :	int=True     ,a__      :	Dict=True     ,a__      :	Optional[Any]=37     ,a__      :	Any="gelu"     ,a__      :	int=10     ,a__      :	Optional[Any]=0.02     ,a__      :	Union[str, Any]=["stage2", "stage3", "stage4"]     ,a__      :	Tuple=3     ,a__      :	Optional[Any]=None     ,)						->       Tuple:




											"""simple docstring"""







											_lowerCAmelCase:List[Any]									=  parent
											_lowerCAmelCase:List[Any]									=  batch_size
											_lowerCAmelCase:Dict									=  image_size
											_lowerCAmelCase:int									=  num_channels
											_lowerCAmelCase:str									=  num_stages
											_lowerCAmelCase:Optional[Any]									=  hidden_sizes
											_lowerCAmelCase:List[Any]									=  depths
											_lowerCAmelCase:int									=  is_training
											_lowerCAmelCase:Dict									=  use_labels
											_lowerCAmelCase:List[Any]									=  intermediate_size
											_lowerCAmelCase:Any									=  hidden_act
											_lowerCAmelCase:Optional[Any]									=  type_sequence_label_size
											_lowerCAmelCase:Optional[Any]									=  initializer_range
											_lowerCAmelCase:Optional[Any]									=  out_features
											_lowerCAmelCase:str									=  num_labels
											_lowerCAmelCase:List[str]									=  scope
											_lowerCAmelCase:Dict									=  num_stages





						def       __UpperCamelCase     (							self      :	str)						->       Dict:




											"""simple docstring"""







											_lowerCAmelCase:Optional[Any]									=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

											_lowerCAmelCase:Tuple									=  None
											if self.use_labels:
																_lowerCAmelCase:str									=  ids_tensor([self.batch_size]     ,self.type_sequence_label_size)

											_lowerCAmelCase:List[str]									=  self.get_config()

											return config, pixel_values, labels





						def       __UpperCamelCase     (							self      :	str)						->       List[str]:




											"""simple docstring"""







											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      :	Any)						->       Optional[Any]:




											"""simple docstring"""







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





						def       __UpperCamelCase     (							self      :	Dict     ,a__      :	Any     ,a__      :	Optional[int]     ,a__      :	Union[str, Any])						->       Optional[Any]:




											"""simple docstring"""







											_lowerCAmelCase:Union[str, Any]									=  UperNetForSemanticSegmentation(config=a__)
											model.to(a__)
											model.eval()
											_lowerCAmelCase:Optional[Any]									=  model(a__)
											self.parent.assertEqual(
											    result.logits.shape     ,(self.batch_size, self.num_labels, self.image_size, self.image_size))







						def       __UpperCamelCase     (							self      :	Union[str, Any])						->       Optional[int]:




											"""simple docstring"""







											_lowerCAmelCase:List[Any]									=  self.prepare_config_and_inputs()
											(
											    _lowerCAmelCase   
											):Tuple									=  config_and_inputs
											_lowerCAmelCase:Dict									=  {'''pixel_values''': pixel_values}
											return config, inputs_dict




@require_torch
class   a__			(       UpperCamelCase_					,	UpperCamelCase_					,	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      :	Union[str, Any])						->       Union[str, Any]:




											"""simple docstring"""







											_lowerCAmelCase:int									=  UperNetModelTester(self)
											_lowerCAmelCase:int									=  ConfigTester(self     ,config_class=a__     ,has_text_modality=a__     ,hidden_size=37)





						def       __UpperCamelCase     (							self      :	Tuple)						->       Any:




											"""simple docstring"""







											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      :	Tuple)						->       List[Any]:




											"""simple docstring"""







											return





						def       __UpperCamelCase     (							self      :	str)						->       List[str]:




											"""simple docstring"""







											_lowerCAmelCase:Optional[int]									=  self.model_tester.prepare_config_and_inputs_for_common()

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

																_lowerCAmelCase:Dict									=  ['''pixel_values''']
																self.assertListEqual(arg_names[:1]     ,a__)





						def       __UpperCamelCase     (							self      :	List[Any])						->       List[Any]:




											"""simple docstring"""







											_lowerCAmelCase:Tuple									=  self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_semantic_segmentation(*a__)





						@unittest.skip(reason='''UperNet does not use inputs_embeds''')
						def       __UpperCamelCase     (							self      :	Tuple)						->       List[Any]:




											"""simple docstring"""







											pass





						@unittest.skip(reason='''UperNet does not support input and output embeddings''')
						def       __UpperCamelCase     (							self      :	int)						->       Any:




											"""simple docstring"""







											pass





						@unittest.skip(reason='''UperNet does not have a base model''')
						def       __UpperCamelCase     (							self      :	str)						->       List[str]:




											"""simple docstring"""







											pass





						@unittest.skip(reason='''UperNet does not have a base model''')
						def       __UpperCamelCase     (							self      :	Tuple)						->       Optional[int]:




											"""simple docstring"""







											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      :	Optional[int])						->       Union[str, Any]:




											"""simple docstring"""







											pass





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




											"""simple docstring"""







											pass





						def       __UpperCamelCase     (							self      :	Optional[int])						->       str:




											"""simple docstring"""







											def check_hidden_states_output(a__      :	List[str]     ,a__      :	int     ,a__      :	str):
																_lowerCAmelCase:Tuple									=  model_class(a__)
																model.to(a__)
																model.eval()

																with torch.no_grad():
																					_lowerCAmelCase:int									=  model(**self._prepare_for_class(a__     ,a__))

																_lowerCAmelCase:Any									=  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																_lowerCAmelCase:int									=  self.model_tester.num_stages
																self.assertEqual(len(a__)     ,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:Dict									=  self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																_lowerCAmelCase:Any									=  True
																check_hidden_states_output(a__     ,a__     ,a__)

																# check that output_hidden_states also work using config
																del inputs_dict["output_hidden_states"]
																_lowerCAmelCase:Union[str, Any]									=  True

																check_hidden_states_output(a__     ,a__     ,a__)





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




											"""simple docstring"""







											_lowerCAmelCase:Optional[int]									=  self.model_tester.prepare_config_and_inputs_for_common()

											_lowerCAmelCase:str									=  _config_zero_init(a__)
											_lowerCAmelCase:List[str]									=  _config_zero_init(configs_no_init.backbone_config)
											for model_class in self.all_model_classes:
																_lowerCAmelCase:int									=  model_class(config=a__)
																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      :	Union[str, Any])						->       str:




											"""simple docstring"""







											pass







						@slow
						def       __UpperCamelCase     (							self      :	int)						->       Any:




											"""simple docstring"""







											for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																_lowerCAmelCase:List[str]									=  UperNetForSemanticSegmentation.from_pretrained(a__)
																self.assertIsNotNone(a__)





def 							UpperCAmelCase		(		):
					_lowerCAmelCase:str									=  hf_hub_download(
					    repo_id='''hf-internal-testing/fixtures_ade20k'''      ,       repo_type='''dataset'''      ,       filename='''ADE_val_00000001.jpg'''    )
					_lowerCAmelCase:int									=  Image.open(snake_case    ).convert('''RGB'''    )
					return image





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





						def       __UpperCamelCase     (							self      :	List[str])						->       Any:




											"""simple docstring"""







											_lowerCAmelCase:Dict									=  AutoImageProcessor.from_pretrained('''openmmlab/upernet-swin-tiny''')
											_lowerCAmelCase:Optional[Any]									=  UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-swin-tiny''').to(a__)

											_lowerCAmelCase:Union[str, Any]									=  prepare_img()
											_lowerCAmelCase:Tuple									=  processor(images=a__     ,return_tensors='''pt''').to(a__)

											with torch.no_grad():
																_lowerCAmelCase:Any									=  model(**a__)

											_lowerCAmelCase:Dict									=  torch.Size((1, model.config.num_labels, 512, 512))
											self.assertEqual(outputs.logits.shape     ,a__)

											_lowerCAmelCase:int									=  torch.tensor(
											    [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]]).to(a__)
											self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]     ,a__     ,atol=1E-4))







						def       __UpperCamelCase     (							self      :	Dict)						->       List[str]:




											"""simple docstring"""







											_lowerCAmelCase:str									=  AutoImageProcessor.from_pretrained('''openmmlab/upernet-convnext-tiny''')
											_lowerCAmelCase:Optional[int]									=  UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-convnext-tiny''').to(a__)

											_lowerCAmelCase:Union[str, Any]									=  prepare_img()
											_lowerCAmelCase:Dict									=  processor(images=a__     ,return_tensors='''pt''').to(a__)

											with torch.no_grad():
																_lowerCAmelCase:Optional[int]									=  model(**a__)

											_lowerCAmelCase:Union[str, Any]									=  torch.Size((1, model.config.num_labels, 512, 512))
											self.assertEqual(outputs.logits.shape     ,a__)

											_lowerCAmelCase:Optional[Any]									=  torch.tensor(
											    [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]]).to(a__)
											self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]     ,a__     ,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 							UpperCAmelCase		(		):
					_lowerCAmelCase:Optional[int]									=  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=snake_case      ,       default=1      ,       help='''Number of TPU cores to use (1 or 8).'''    )

					# positional
					parser.add_argument(
					    '''training_script'''      ,       type=snake_case      ,       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=snake_case    )

					return parser.parse_args()





def 							UpperCAmelCase		(		):
					_lowerCAmelCase:int									=  parse_args()

					# Import training_script as a module.
					_lowerCAmelCase:Optional[int]									=  Path(args.training_script    )
					sys.path.append(str(script_fpath.parent.resolve()    )    )
					_lowerCAmelCase:Optional[int]									=  script_fpath.stem
					_lowerCAmelCase:Any									=  importlib.import_module(snake_case    )

					# Patch sys.argv
					_lowerCAmelCase:Dict									=  [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()






