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




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



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       parent
												SCREAMING_SNAKE_CASE    			=       batch_size
												SCREAMING_SNAKE_CASE    			=       image_size
												SCREAMING_SNAKE_CASE    			=       num_channels
												SCREAMING_SNAKE_CASE    			=       num_stages
												SCREAMING_SNAKE_CASE    			=       hidden_sizes
												SCREAMING_SNAKE_CASE    			=       depths
												SCREAMING_SNAKE_CASE    			=       is_training
												SCREAMING_SNAKE_CASE    			=       use_labels
												SCREAMING_SNAKE_CASE    			=       intermediate_size
												SCREAMING_SNAKE_CASE    			=       hidden_act
												SCREAMING_SNAKE_CASE    			=       type_sequence_label_size
												SCREAMING_SNAKE_CASE    			=       initializer_range
												SCREAMING_SNAKE_CASE    			=       out_features
												SCREAMING_SNAKE_CASE    			=       num_labels
												SCREAMING_SNAKE_CASE    			=       scope
												SCREAMING_SNAKE_CASE    			=       num_stages




							def 	__snake_case(						self   :					Optional[Any]      )					->				Optional[int]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

												SCREAMING_SNAKE_CASE    			=       self.get_config()

												return config, pixel_values, labels




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



												'''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 	__snake_case(						self   :					List[Any]      )					->				Any:



												'''simple docstring'''

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




							def 	__snake_case(						self   :					str	,							_UpperCamelCase   :					Tuple	,							_UpperCamelCase   :					List[Any]	,							_UpperCamelCase   :					Optional[int]      )					->				Optional[int]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       UperNetForSemanticSegmentation(config=_UpperCamelCase      )
												model.to(_UpperCamelCase      )
												model.eval()
												SCREAMING_SNAKE_CASE    			=       model(_UpperCamelCase      )
												self.parent.assertEqual(
												    result.logits.shape	,							(self.batch_size, self.num_labels, self.image_size, self.image_size)      )





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



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.prepare_config_and_inputs()
												(
												    SCREAMING_SNAKE_CASE							
												)    			=       config_and_inputs
												SCREAMING_SNAKE_CASE    			=       {"pixel_values": pixel_values}
												return config, inputs_dict



@require_torch
class 						lowercase    (     lowercase__				,       lowercase__				,       unittest.TestCase     ):
							lowercase__						:       Dict								=							(UperNetForSemanticSegmentation,) if is_torch_available() else ()
							lowercase__						:       int								=							{"""image-segmentation""": UperNetForSemanticSegmentation} if is_torch_available() else {}
							lowercase__						:       Union[str, Any]								=							False
							lowercase__						:       int								=							False
							lowercase__						:       Optional[Any]								=							False
							lowercase__						:       Any								=							False
							lowercase__						:       Optional[int]								=							False
							lowercase__						:       Dict								=							False




							def 	__snake_case(						self   :					Any      )					->				Optional[int]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       UperNetModelTester(self      )
												SCREAMING_SNAKE_CASE    			=       ConfigTester(self	,							config_class=_UpperCamelCase	,							has_text_modality=_UpperCamelCase	,							hidden_size=37      )




							def 	__snake_case(						self   :					Dict      )					->				Union[str, 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 	__snake_case(						self   :					List[str]      )					->				Optional[Any]:



												'''simple docstring'''

												return




							def 	__snake_case(						self   :					Dict      )					->				Dict:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs_for_common()

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

																	SCREAMING_SNAKE_CASE    			=       ["pixel_values"]
																	self.assertListEqual(arg_names[:1]	,							_UpperCamelCase      )




							def 	__snake_case(						self   :					List[Any]      )					->				Optional[Any]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_semantic_segmentation(*_UpperCamelCase      )




							@unittest.skip(reason="UperNet does not use inputs_embeds"      )
							def 	__snake_case(						self   :					Optional[Any]      )					->				str:



												'''simple docstring'''

												pass




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



												'''simple docstring'''

												pass




							@unittest.skip(reason="UperNet does not have a base model"      )
							def 	__snake_case(						self   :					Union[str, Any]      )					->				List[Any]:



												'''simple docstring'''

												pass




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



												'''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 	__snake_case(						self   :					Dict      )					->				List[str]:



												'''simple docstring'''

												pass




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



												'''simple docstring'''

												pass




							def 	__snake_case(						self   :					List[str]      )					->				Tuple:



												'''simple docstring'''

												def check_hidden_states_output(_UpperCamelCase   :					Tuple	,							_UpperCamelCase   :					Any	,							_UpperCamelCase   :					Optional[Any]      ):
																	SCREAMING_SNAKE_CASE    			=       model_class(_UpperCamelCase      )
																	model.to(_UpperCamelCase      )
																	model.eval()

																	with torch.no_grad():
																						SCREAMING_SNAKE_CASE    			=       model(**self._prepare_for_class(_UpperCamelCase	,							_UpperCamelCase      )      )

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

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

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																	SCREAMING_SNAKE_CASE    			=       True
																	check_hidden_states_output(_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase      )

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

																	check_hidden_states_output(_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase      )




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



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       self.model_tester.prepare_config_and_inputs_for_common()

												SCREAMING_SNAKE_CASE    			=       _config_zero_init(_UpperCamelCase      )
												SCREAMING_SNAKE_CASE    			=       _config_zero_init(configs_no_init.backbone_config      )
												for model_class in self.all_model_classes:
																	SCREAMING_SNAKE_CASE    			=       model_class(config=_UpperCamelCase      )
																	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   :					List[str]      )					->				List[str]:



												'''simple docstring'''

												pass





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



												'''simple docstring'''

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





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


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




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



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny"      )
												SCREAMING_SNAKE_CASE    			=       UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny"      ).to(_UpperCamelCase      )

												SCREAMING_SNAKE_CASE    			=       prepare_img()
												SCREAMING_SNAKE_CASE    			=       processor(images=_UpperCamelCase	,							return_tensors="pt"      ).to(_UpperCamelCase      )

												with torch.no_grad():
																	SCREAMING_SNAKE_CASE    			=       model(**_UpperCamelCase      )

												SCREAMING_SNAKE_CASE    			=       torch.Size((1, model.config.num_labels, 512, 512)      )
												self.assertEqual(outputs.logits.shape	,							_UpperCamelCase      )

												SCREAMING_SNAKE_CASE    			=       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(_UpperCamelCase      )
												self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]	,							_UpperCamelCase	,							atol=1e-4      )      )





							def 	__snake_case(						self   :					Tuple      )					->				Optional[int]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny"      )
												SCREAMING_SNAKE_CASE    			=       UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny"      ).to(_UpperCamelCase      )

												SCREAMING_SNAKE_CASE    			=       prepare_img()
												SCREAMING_SNAKE_CASE    			=       processor(images=_UpperCamelCase	,							return_tensors="pt"      ).to(_UpperCamelCase      )

												with torch.no_grad():
																	SCREAMING_SNAKE_CASE    			=       model(**_UpperCamelCase      )

												SCREAMING_SNAKE_CASE    			=       torch.Size((1, model.config.num_labels, 512, 512)      )
												self.assertEqual(outputs.logits.shape	,							_UpperCamelCase      )

												SCREAMING_SNAKE_CASE    			=       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(_UpperCamelCase      )
												self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]	,							_UpperCamelCase	,							atol=1e-4      )      )





import inspect
import os
import unittest

import torch

import accelerate
from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_multi_gpu
from accelerate.utils import patch_environment


class 						lowercase    (     unittest.TestCase     ):




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



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       inspect.getfile(accelerate.test_utils      )
												SCREAMING_SNAKE_CASE    			=       os.path.sep.join(mod_file.split(os.path.sep      )[:-1] + ["scripts", "test_script.py"]      )
												SCREAMING_SNAKE_CASE    			=       os.path.sep.join(
												    mod_file.split(os.path.sep      )[:-1] + ["scripts", "test_distributed_data_loop.py"]      )
												SCREAMING_SNAKE_CASE    			=       os.path.sep.join(mod_file.split(os.path.sep      )[:-1] + ["scripts", "test_ops.py"]      )




							@require_multi_gpu
							def 	__snake_case(						self   :					Optional[int]      )					->				Any:



												'''simple docstring'''

												print(F"Found {torch.cuda.device_count()} devices."      )
												SCREAMING_SNAKE_CASE    			=       ["torchrun", F"--nproc_per_node={torch.cuda.device_count()}", self.test_file_path]
												with patch_environment(omp_num_threads=1      ):
																	execute_subprocess_async(_UpperCamelCase	,							env=os.environ.copy()      )




							@require_multi_gpu
							def 	__snake_case(						self   :					List[Any]      )					->				int:



												'''simple docstring'''

												print(F"Found {torch.cuda.device_count()} devices."      )
												SCREAMING_SNAKE_CASE    			=       ["torchrun", F"--nproc_per_node={torch.cuda.device_count()}", self.operation_file_path]
												print(F"Command: {cmd}"      )
												with patch_environment(omp_num_threads=1      ):
																	execute_subprocess_async(_UpperCamelCase	,							env=os.environ.copy()      )




							@require_multi_gpu
							def 	__snake_case(						self   :					int      )					->				Optional[int]:



												'''simple docstring'''

												SCREAMING_SNAKE_CASE    			=       ["torchrun", F"--nproc_per_node={torch.cuda.device_count()}", inspect.getfile(self.__class__      )]
												with patch_environment(omp_num_threads=1      ):
																	execute_subprocess_async(_UpperCamelCase	,							env=os.environ.copy()      )





							@require_multi_gpu
							def 	__snake_case(						self   :					int      )					->				int:



												'''simple docstring'''

												print(F"Found {torch.cuda.device_count()} devices, using 2 devices only"      )
												SCREAMING_SNAKE_CASE    			=       ["torchrun", F"--nproc_per_node={torch.cuda.device_count()}", self.data_loop_file_path]
												with patch_environment(omp_num_threads=1	,							cuda_visible_devices="0,1"      ):
																	execute_subprocess_async(_UpperCamelCase	,							env=os.environ.copy()      )


if __name__ == "__main__":
							_lowerCamelCase					:   str							=		Accelerator()
							_lowerCamelCase					:   List[str]							=		(accelerator.state.process_index + 2, 10)
							_lowerCamelCase					:   str							=		torch.randint(0, 10, shape).to(accelerator.device)

							_lowerCamelCase					:   Optional[Any]							=		''''''

							_lowerCamelCase					:   str							=		accelerator.pad_across_processes(tensor)
							if tensora.shape[0] != accelerator.state.num_processes + 1:
														error_msg += f"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
							if not torch.equal(tensora[: accelerator.state.process_index + 2], tensor):
														error_msg += "Tensors have different values."
							if not torch.all(tensora[accelerator.state.process_index + 2 :] == 0):
														error_msg += "Padding was not done with the right value (0)."

							_lowerCamelCase					:   Any							=		accelerator.pad_across_processes(tensor, pad_first=True)
							if tensora.shape[0] != accelerator.state.num_processes + 1:
														error_msg += f"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
							_lowerCamelCase					:   int							=		accelerator.state.num_processes - accelerator.state.process_index - 1
							if not torch.equal(tensora[index:], tensor):
														error_msg += "Tensors have different values."
							if not torch.all(tensora[:index] == 0):
														error_msg += "Padding was not done with the right value (0)."

							# Raise error at the end to make sure we don't stop at the first failure.
							if len(error_msg) > 0:
														raise ValueError(error_msg)




