





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							,					_A							,					_A=1_3							,					_A=3_2							,					_A=3							,					_A=4							,					_A=[1_0, 2_0, 3_0, 4_0]							,					_A=[2, 2, 3, 2]							,					_A=True							,					_A=True							,					_A=3_7							,					_A="gelu"							,					_A=1_0							,					_A=0.02							,					_A=["stage2", "stage3", "stage4"]							,					_A=3							,					_A=None							,					):






				'''simple docstring'''







				UpperCAmelCase   		=				parent
				UpperCAmelCase   		=				batch_size
				UpperCAmelCase   		=				image_size
				UpperCAmelCase   		=				num_channels
				UpperCAmelCase   		=				num_stages
				UpperCAmelCase   		=				hidden_sizes
				UpperCAmelCase   		=				depths
				UpperCAmelCase   		=				is_training
				UpperCAmelCase   		=				use_labels
				UpperCAmelCase   		=				intermediate_size
				UpperCAmelCase   		=				hidden_act
				UpperCAmelCase   		=				type_sequence_label_size
				UpperCAmelCase   		=				initializer_range
				UpperCAmelCase   		=				out_features
				UpperCAmelCase   		=				num_labels
				UpperCAmelCase   		=				scope
				UpperCAmelCase   		=				num_stages
	def      _lowercase				(    self		):






				'''simple docstring'''







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

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

				UpperCAmelCase   		=				self.get_config()

				return config, pixel_values, labels
	def      _lowercase				(    self		):






				'''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      _lowercase				(    self		):






				'''simple docstring'''







				return UperNetConfig(
				    backbone_config=self.get_backbone_config()							,					hidden_size=5_1_2							,					pool_scales=[1, 2, 3, 6]							,					use_auxiliary_head=_A							,					auxiliary_loss_weight=0.4							,					auxiliary_in_channels=4_0							,					auxiliary_channels=2_5_6							,					auxiliary_num_convs=1							,					auxiliary_concat_input=_A							,					loss_ignore_index=2_5_5							,					num_labels=self.num_labels							,					)
	def      _lowercase				(    self							,					_A							,					_A							,					_A		):






				'''simple docstring'''







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







	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.prepare_config_and_inputs()
				(
				    (
				    UpperCAmelCase
				)		,			(
				    UpperCAmelCase
				)		,			(
				    UpperCAmelCase
				)		,			
				)   		=				config_and_inputs
				UpperCAmelCase   		=				{'''pixel_values''': pixel_values}
				return config, inputs_dict

@require_torch
class 	A_       (a_     ,							a_     ,							unittest.TestCase			):
	UpperCAmelCase__					=			(UperNetForSemanticSegmentation,) if is_torch_available() else ()
	UpperCAmelCase__					=			{'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
	UpperCAmelCase__					=			False
	UpperCAmelCase__					=			False
	UpperCAmelCase__					=			False
	UpperCAmelCase__					=			False
	UpperCAmelCase__					=			False
	UpperCAmelCase__					=			False
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				UperNetModelTester(self		)
				UpperCAmelCase   		=				ConfigTester(self							,					config_class=_A							,					has_text_modality=_A							,					hidden_size=3_7		)
	def      _lowercase				(    self		):






				'''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      _lowercase				(    self		):






				'''simple docstring'''







				return
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase		,			UpperCAmelCase   		=				self.model_tester.prepare_config_and_inputs_for_common()

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

							UpperCAmelCase   		=				['''pixel_values''']
							self.assertListEqual(arg_names[:1]							,					_A		)
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				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      _lowercase				(    self		):






				'''simple docstring'''







				pass
	@unittest.skip(reason='''UperNet does not support input and output embeddings'''		)
	def      _lowercase				(    self		):






				'''simple docstring'''







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






				'''simple docstring'''







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






				'''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      _lowercase				(    self		):






				'''simple docstring'''







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






				'''simple docstring'''







				pass
	def      _lowercase				(    self		):






				'''simple docstring'''







				def check_hidden_states_output(_A							,					_A							,					_A		):
							UpperCAmelCase   		=				model_class(_A		)
							model.to(_A		)
							model.eval()

							with torch.no_grad():
										UpperCAmelCase   		=				model(**self._prepare_for_class(_A							,					_A		)		)

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

							UpperCAmelCase   		=				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]							,					)

				UpperCAmelCase		,			UpperCAmelCase   		=				self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
							UpperCAmelCase   		=				True
							check_hidden_states_output(_A							,					_A							,					_A		)

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

							check_hidden_states_output(_A							,					_A							,					_A		)
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase		,			UpperCAmelCase   		=				self.model_tester.prepare_config_and_inputs_for_common()

				UpperCAmelCase   		=				_config_zero_init(_A		)
				UpperCAmelCase   		=				_config_zero_init(configs_no_init.backbone_config		)
				for model_class in self.all_model_classes:
							UpperCAmelCase   		=				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      _lowercase				(    self		):






				'''simple docstring'''







				pass







	@slow
	def      _lowercase				(    self		):






				'''simple docstring'''







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





def 						__SCREAMING_SNAKE_CASE  (				)  ->				Optional[int]:







			'''simple docstring'''

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






@require_torch
@require_vision
@slow
class 	A_       (unittest.TestCase			):
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				AutoImageProcessor.from_pretrained('''openmmlab/upernet-swin-tiny'''		)
				UpperCAmelCase   		=				UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-swin-tiny'''		).to(_A		)

				UpperCAmelCase   		=				prepare_img()
				UpperCAmelCase   		=				processor(images=_A							,					return_tensors='''pt'''		).to(_A		)

				with torch.no_grad():
							UpperCAmelCase   		=				model(**_A		)

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

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







	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				AutoImageProcessor.from_pretrained('''openmmlab/upernet-convnext-tiny'''		)
				UpperCAmelCase   		=				UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-convnext-tiny'''		).to(_A		)

				UpperCAmelCase   		=				prepare_img()
				UpperCAmelCase   		=				processor(images=_A							,					return_tensors='''pt'''		).to(_A		)

				with torch.no_grad():
							UpperCAmelCase   		=				model(**_A		)

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

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










import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np

class 	A_       :
	def      _lowercase				(    self		):






				'''simple docstring'''







				torch.manual_seed(0		)
				UpperCAmelCase   		=				TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				UNetaDConditionModel(
				    sample_size=3_2							,					layers_per_block=1							,					block_out_channels=[3_2, 6_4]							,					down_block_types=[
				        '''ResnetDownsampleBlock2D''',
				        '''SimpleCrossAttnDownBlock2D''',
				    ]							,					mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,					up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,					in_channels=3							,					out_channels=6							,					cross_attention_dim=3_2							,					encoder_hid_dim=3_2							,					attention_head_dim=8							,					addition_embed_type='''text'''							,					addition_embed_type_num_heads=2							,					cross_attention_norm='''group_norm'''							,					resnet_time_scale_shift='''scale_shift'''							,					act_fn='''gelu'''							,					)
				unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				torch.manual_seed(0		)
				UpperCAmelCase   		=				DDPMScheduler(
				    num_train_timesteps=1_0_0_0							,					beta_schedule='''squaredcos_cap_v2'''							,					beta_start=0.00_01							,					beta_end=0.02							,					thresholding=_A							,					dynamic_thresholding_ratio=0.95							,					sample_max_value=1.0							,					prediction_type='''epsilon'''							,					variance_type='''learned_range'''							,					)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				IFWatermarker()

				return {
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "unet": unet,
				    "scheduler": scheduler,
				    "watermarker": watermarker,
				    "safety_checker": None,
				    "feature_extractor": None,
				}
	def      _lowercase				(    self		):






				'''simple docstring'''







				torch.manual_seed(0		)
				UpperCAmelCase   		=				TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				UNetaDConditionModel(
				    sample_size=3_2							,					layers_per_block=[1, 2]							,					block_out_channels=[3_2, 6_4]							,					down_block_types=[
				        '''ResnetDownsampleBlock2D''',
				        '''SimpleCrossAttnDownBlock2D''',
				    ]							,					mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,					up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,					in_channels=6							,					out_channels=6							,					cross_attention_dim=3_2							,					encoder_hid_dim=3_2							,					attention_head_dim=8							,					addition_embed_type='''text'''							,					addition_embed_type_num_heads=2							,					cross_attention_norm='''group_norm'''							,					resnet_time_scale_shift='''scale_shift'''							,					act_fn='''gelu'''							,					class_embed_type='''timestep'''							,					mid_block_scale_factor=1.4_14							,					time_embedding_act_fn='''gelu'''							,					time_embedding_dim=3_2							,					)
				unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				torch.manual_seed(0		)
				UpperCAmelCase   		=				DDPMScheduler(
				    num_train_timesteps=1_0_0_0							,					beta_schedule='''squaredcos_cap_v2'''							,					beta_start=0.00_01							,					beta_end=0.02							,					thresholding=_A							,					dynamic_thresholding_ratio=0.95							,					sample_max_value=1.0							,					prediction_type='''epsilon'''							,					variance_type='''learned_range'''							,					)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				DDPMScheduler(
				    num_train_timesteps=1_0_0_0							,					beta_schedule='''squaredcos_cap_v2'''							,					beta_start=0.00_01							,					beta_end=0.02							,					)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				IFWatermarker()

				return {
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "unet": unet,
				    "scheduler": scheduler,
				    "image_noising_scheduler": image_noising_scheduler,
				    "watermarker": watermarker,
				    "safety_checker": None,
				    "feature_extractor": None,
				}
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.get_dummy_components()
				UpperCAmelCase   		=				self.pipeline_class(**_A		)
				pipe.to(_A		)
				pipe.set_progress_bar_config(disable=_A		)

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)

				UpperCAmelCase   		=				inputs['''prompt''']
				UpperCAmelCase   		=				inputs['''generator''']
				UpperCAmelCase   		=				inputs['''num_inference_steps''']
				UpperCAmelCase   		=				inputs['''output_type''']

				if "image" in inputs:
							UpperCAmelCase   		=				inputs['''image''']
				else:
							UpperCAmelCase   		=				None

				if "mask_image" in inputs:
							UpperCAmelCase   		=				inputs['''mask_image''']
				else:
							UpperCAmelCase   		=				None

				if "original_image" in inputs:
							UpperCAmelCase   		=				inputs['''original_image''']
				else:
							UpperCAmelCase   		=				None

				UpperCAmelCase		,			UpperCAmelCase   		=				pipe.encode_prompt(_A		)

				# inputs with prompt converted to embeddings
				UpperCAmelCase   		=				{
				    '''prompt_embeds''': prompt_embeds,
				    '''negative_prompt_embeds''': negative_prompt_embeds,
				    '''generator''': generator,
				    '''num_inference_steps''': num_inference_steps,
				    '''output_type''': output_type,
				}

				if image is not None:
							UpperCAmelCase   		=				image

				if mask_image is not None:
							UpperCAmelCase   		=				mask_image

				if original_image is not None:
							UpperCAmelCase   		=				original_image

				# set all optional components to None
				for optional_component in pipe._optional_components:
							setattr(_A							,					_A							,					_A		)

				UpperCAmelCase   		=				pipe(**_A		)[0]

				with tempfile.TemporaryDirectory() as tmpdir:
							pipe.save_pretrained(_A		)
							UpperCAmelCase   		=				self.pipeline_class.from_pretrained(_A		)
							pipe_loaded.to(_A		)
							pipe_loaded.set_progress_bar_config(disable=_A		)

				pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				for optional_component in pipe._optional_components:
							self.assertTrue(
							    getattr(_A							,					_A		) is None							,					F"""`{optional_component}` did not stay set to None after loading."""							,					)

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)

				UpperCAmelCase   		=				inputs['''generator''']
				UpperCAmelCase   		=				inputs['''num_inference_steps''']
				UpperCAmelCase   		=				inputs['''output_type''']

				# inputs with prompt converted to embeddings
				UpperCAmelCase   		=				{
				    '''prompt_embeds''': prompt_embeds,
				    '''negative_prompt_embeds''': negative_prompt_embeds,
				    '''generator''': generator,
				    '''num_inference_steps''': num_inference_steps,
				    '''output_type''': output_type,
				}

				if image is not None:
							UpperCAmelCase   		=				image

				if mask_image is not None:
							UpperCAmelCase   		=				mask_image

				if original_image is not None:
							UpperCAmelCase   		=				original_image

				UpperCAmelCase   		=				pipe_loaded(**_A		)[0]

				UpperCAmelCase   		=				np.abs(to_np(_A		) - to_np(_A		)		).max()
				self.assertLess(_A							,					1E-4		)







	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.get_dummy_components()
				UpperCAmelCase   		=				self.pipeline_class(**_A		)
				pipe.to(_A		)
				pipe.set_progress_bar_config(disable=_A		)

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)
				UpperCAmelCase   		=				pipe(**_A		)[0]

				with tempfile.TemporaryDirectory() as tmpdir:
							pipe.save_pretrained(_A		)
							UpperCAmelCase   		=				self.pipeline_class.from_pretrained(_A		)
							pipe_loaded.to(_A		)
							pipe_loaded.set_progress_bar_config(disable=_A		)

				pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)
				UpperCAmelCase   		=				pipe_loaded(**_A		)[0]

				UpperCAmelCase   		=				np.abs(to_np(_A		) - to_np(_A		)		).max()
				self.assertLess(_A							,					1E-4		)



