
'''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    lowerCAmelCase       :


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

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


 def    _A	(    self  : Any    ):
  '''simple docstring'''

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

  lowerCAmelCase__     :   Optional[Any]     		=	None
  if self.use_labels:
   lowerCAmelCase__     :   Any     		=	ids_tensor([self.batch_size]			,							self.type_sequence_label_size    )

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

  return config, pixel_values, labels


 def    _A	(    self  : List[Any]    ):
  '''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    _A	(    self  : Optional[int]    ):
  '''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    _A	(    self  : Optional[Any]			,							a__  : List[Any]			,							a__  : List[Any]			,							a__  : int    ):
  '''simple docstring'''

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


 def    _A	(    self  : Optional[Any]    ):
  '''simple docstring'''

  lowerCAmelCase__     :   Optional[int]     		=	self.prepare_config_and_inputs()
  (
      (
      lowerCAmelCase__
  )		,     (
      lowerCAmelCase__
  )		,     (
      lowerCAmelCase__
  )		,     
  )     :   Any     		=	config_and_inputs
  lowerCAmelCase__     :   Dict     		=	{"pixel_values": pixel_values}
  return config, inputs_dict




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


 def    _A	(    self  : List[str]    ):
  '''simple docstring'''

  lowerCAmelCase__     :   Tuple     		=	UperNetModelTester(self    )
  lowerCAmelCase__     :   List[str]     		=	ConfigTester(self			,							config_class=a__			,							has_text_modality=a__			,							hidden_size=37    )


 def    _A	(    self  : Dict    ):
  '''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    _A	(    self  : Any    ):
  '''simple docstring'''

  return


 def    _A	(    self  : Optional[int]    ):
  '''simple docstring'''

  lowerCAmelCase__		,     lowerCAmelCase__     :   Optional[Any]     		=	self.model_tester.prepare_config_and_inputs_for_common()

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

   lowerCAmelCase__     :   Dict     		=	["pixel_values"]
   self.assertListEqual(arg_names[:1]			,							a__    )


 def    _A	(    self  : Union[str, Any]    ):
  '''simple docstring'''

  lowerCAmelCase__     :   Union[str, Any]     		=	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    _A	(    self  : Any    ):
  '''simple docstring'''

  pass


 @unittest.skip(reason="UperNet does not support input and output embeddings"    )
 def    _A	(    self  : int    ):
  '''simple docstring'''

  pass


 @unittest.skip(reason="UperNet does not have a base model"    )
 def    _A	(    self  : Optional[Any]    ):
  '''simple docstring'''

  pass


 @unittest.skip(reason="UperNet does not have a base model"    )
 def    _A	(    self  : str    ):
  '''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    _A	(    self  : int    ):
  '''simple docstring'''

  pass


 @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."    )
 def    _A	(    self  : Tuple    ):
  '''simple docstring'''

  pass


 def    _A	(    self  : Tuple    ):
  '''simple docstring'''

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

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

   lowerCAmelCase__     :   Union[str, Any]     		=	outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

   lowerCAmelCase__     :   Dict     		=	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__		,     lowerCAmelCase__     :   Optional[int]     		=	self.model_tester.prepare_config_and_inputs_for_common()

  for model_class in self.all_model_classes:
   lowerCAmelCase__     :   Union[str, 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__     :   List[Any]     		=	True

   check_hidden_states_output(a__			,							a__			,							a__    )


 def    _A	(    self  : List[Any]    ):
  '''simple docstring'''

  lowerCAmelCase__		,     lowerCAmelCase__     :   int     		=	self.model_tester.prepare_config_and_inputs_for_common()

  lowerCAmelCase__     :   Union[str, Any]     		=	_config_zero_init(a__    )
  lowerCAmelCase__     :   Optional[Any]     		=	_config_zero_init(configs_no_init.backbone_config    )
  for model_class in self.all_model_classes:
   lowerCAmelCase__     :   Union[str, Any]     		=	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    _A	(    self  : Any    ):
  '''simple docstring'''

  pass


 @slow
 def    _A	(    self  : Any    ):
  '''simple docstring'''

  for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   lowerCAmelCase__     :   Optional[int]     		=	UperNetForSemanticSegmentation.from_pretrained(a__    )
   self.assertIsNotNone(a__    )



def        UpperCAmelCase_ (  ):







 """simple docstring"""


 lowerCAmelCase__     :   List[Any]     		=	hf_hub_download(
     repo_id="hf-internal-testing/fixtures_ade20k"      ,			repo_type="dataset"      ,			filename="ADE_val_00000001.jpg"			)
 lowerCAmelCase__     :   Tuple     		=	Image.open(lowerCamelCase_			).convert("RGB"			)
 return image







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


 def    _A	(    self  : Dict    ):
  '''simple docstring'''

  lowerCAmelCase__     :   Union[str, Any]     		=	AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny"    )
  lowerCAmelCase__     :   Tuple     		=	UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny"    ).to(a__    )

  lowerCAmelCase__     :   List[str]     		=	prepare_img()
  lowerCAmelCase__     :   int     		=	processor(images=a__			,							return_tensors="pt"    ).to(a__    )

  with torch.no_grad():
   lowerCAmelCase__     :   int     		=	model(**a__    )

  lowerCAmelCase__     :   Union[str, Any]     		=	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    _A	(    self  : Optional[int]    ):
  '''simple docstring'''

  lowerCAmelCase__     :   int     		=	AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny"    )
  lowerCAmelCase__     :   str     		=	UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny"    ).to(a__    )

  lowerCAmelCase__     :   List[Any]     		=	prepare_img()
  lowerCAmelCase__     :   Dict     		=	processor(images=a__			,							return_tensors="pt"    ).to(a__    )

  with torch.no_grad():
   lowerCAmelCase__     :   Dict     		=	model(**a__    )

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

  lowerCAmelCase__     :   Dict     		=	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 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    lowerCAmelCase       :


 def    _A	(    self  : Union[str, Any]    ):
  '''simple docstring'''

  torch.manual_seed(0    )
  lowerCAmelCase__     :   List[str]     		=	TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Optional[int]     		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Any     		=	UNetaDConditionModel(
      sample_size=32			,							layers_per_block=1			,							block_out_channels=[32, 64]			,							down_block_types=[
          "ResnetDownsampleBlock2D",
          "SimpleCrossAttnDownBlock2D",
      ]			,							mid_block_type="UNetMidBlock2DSimpleCrossAttn"			,							up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]			,							in_channels=3			,							out_channels=6			,							cross_attention_dim=32			,							encoder_hid_dim=32			,							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    )
  lowerCAmelCase__     :   Union[str, Any]     		=	DDPMScheduler(
      num_train_timesteps=1000			,							beta_schedule="squaredcos_cap_v2"			,							beta_start=0.0001			,							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    )
  lowerCAmelCase__     :   Tuple     		=	IFWatermarker()

  return {
      "text_encoder": text_encoder,
      "tokenizer": tokenizer,
      "unet": unet,
      "scheduler": scheduler,
      "watermarker": watermarker,
      "safety_checker": None,
      "feature_extractor": None,
  }


 def    _A	(    self  : Dict    ):
  '''simple docstring'''

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Any     		=	TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   List[Any]     		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"    )

  torch.manual_seed(0    )
  lowerCAmelCase__     :   int     		=	UNetaDConditionModel(
      sample_size=32			,							layers_per_block=[1, 2]			,							block_out_channels=[32, 64]			,							down_block_types=[
          "ResnetDownsampleBlock2D",
          "SimpleCrossAttnDownBlock2D",
      ]			,							mid_block_type="UNetMidBlock2DSimpleCrossAttn"			,							up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]			,							in_channels=6			,							out_channels=6			,							cross_attention_dim=32			,							encoder_hid_dim=32			,							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.414			,							time_embedding_act_fn="gelu"			,							time_embedding_dim=32			,							)
  unet.set_attn_processor(AttnAddedKVProcessor()    )  # For reproducibility tests

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Tuple     		=	DDPMScheduler(
      num_train_timesteps=1000			,							beta_schedule="squaredcos_cap_v2"			,							beta_start=0.0001			,							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    )
  lowerCAmelCase__     :   Tuple     		=	DDPMScheduler(
      num_train_timesteps=1000			,							beta_schedule="squaredcos_cap_v2"			,							beta_start=0.0001			,							beta_end=0.02			,							)

  torch.manual_seed(0    )
  lowerCAmelCase__     :   Dict     		=	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    _A	(    self  : Optional[Any]    ):
  '''simple docstring'''

  lowerCAmelCase__     :   Union[str, Any]     		=	self.get_dummy_components()
  lowerCAmelCase__     :   Dict     		=	self.pipeline_class(**a__    )
  pipe.to(a__    )
  pipe.set_progress_bar_config(disable=a__    )

  lowerCAmelCase__     :   Tuple     		=	self.get_dummy_inputs(a__    )

  lowerCAmelCase__     :   str     		=	inputs["prompt"]
  lowerCAmelCase__     :   Optional[Any]     		=	inputs["generator"]
  lowerCAmelCase__     :   str     		=	inputs["num_inference_steps"]
  lowerCAmelCase__     :   int     		=	inputs["output_type"]

  if "image" in inputs:
   lowerCAmelCase__     :   Optional[int]     		=	inputs["image"]
  else:
   lowerCAmelCase__     :   str     		=	None

  if "mask_image" in inputs:
   lowerCAmelCase__     :   List[str]     		=	inputs["mask_image"]
  else:
   lowerCAmelCase__     :   str     		=	None

  if "original_image" in inputs:
   lowerCAmelCase__     :   Optional[int]     		=	inputs["original_image"]
  else:
   lowerCAmelCase__     :   Union[str, Any]     		=	None

  lowerCAmelCase__		,     lowerCAmelCase__     :   Dict     		=	pipe.encode_prompt(a__    )

  # inputs with prompt converted to embeddings
  lowerCAmelCase__     :   Dict     		=	{
      "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:
   lowerCAmelCase__     :   Optional[Any]     		=	image

  if mask_image is not None:
   lowerCAmelCase__     :   Optional[int]     		=	mask_image

  if original_image is not None:
   lowerCAmelCase__     :   Optional[Any]     		=	original_image

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

  lowerCAmelCase__     :   Optional[int]     		=	pipe(**a__    )[0]

  with tempfile.TemporaryDirectory() as tmpdir:
   pipe.save_pretrained(a__    )
   lowerCAmelCase__     :   Tuple     		=	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.'''			,							)

  lowerCAmelCase__     :   Union[str, Any]     		=	self.get_dummy_inputs(a__    )

  lowerCAmelCase__     :   Any     		=	inputs["generator"]
  lowerCAmelCase__     :   Union[str, Any]     		=	inputs["num_inference_steps"]
  lowerCAmelCase__     :   str     		=	inputs["output_type"]

  # inputs with prompt converted to embeddings
  lowerCAmelCase__     :   int     		=	{
      "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:
   lowerCAmelCase__     :   Tuple     		=	image

  if mask_image is not None:
   lowerCAmelCase__     :   int     		=	mask_image

  if original_image is not None:
   lowerCAmelCase__     :   List[Any]     		=	original_image

  lowerCAmelCase__     :   Dict     		=	pipe_loaded(**a__    )[0]

  lowerCAmelCase__     :   int     		=	np.abs(to_np(a__    ) - to_np(a__    )    ).max()
  self.assertLess(a__			,							1e-4    )


 def    _A	(    self  : Any    ):
  '''simple docstring'''

  lowerCAmelCase__     :   str     		=	self.get_dummy_components()
  lowerCAmelCase__     :   int     		=	self.pipeline_class(**a__    )
  pipe.to(a__    )
  pipe.set_progress_bar_config(disable=a__    )

  lowerCAmelCase__     :   Any     		=	self.get_dummy_inputs(a__    )
  lowerCAmelCase__     :   Union[str, Any]     		=	pipe(**a__    )[0]

  with tempfile.TemporaryDirectory() as tmpdir:
   pipe.save_pretrained(a__    )
   lowerCAmelCase__     :   str     		=	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

  lowerCAmelCase__     :   str     		=	self.get_dummy_inputs(a__    )
  lowerCAmelCase__     :   Union[str, Any]     		=	pipe_loaded(**a__    )[0]

  lowerCAmelCase__     :   Union[str, Any]     		=	np.abs(to_np(a__    ) - to_np(a__    )    ).max()
  self.assertLess(a__			,							1e-4    )





