
"""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 			UpperCamelCase      :
     def __init__(     self  ,     snake_case__  ,     snake_case__=13  ,     snake_case__=32  ,     snake_case__=3  ,     snake_case__=4  ,     snake_case__=[10, 20, 30, 40]  ,     snake_case__=[2, 2, 3, 2]  ,     snake_case__=True  ,     snake_case__=True  ,     snake_case__=37  ,     snake_case__="gelu"  ,     snake_case__=10  ,     snake_case__=0.02  ,     snake_case__=["stage2", "stage3", "stage4"]  ,     snake_case__=3  ,     snake_case__=None  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = parent
       _SCREAMING_SNAKE_CASE   :				List[str]							    = batch_size
       _SCREAMING_SNAKE_CASE   :				Tuple							    = image_size
       _SCREAMING_SNAKE_CASE   :				List[str]							    = num_channels
       _SCREAMING_SNAKE_CASE   :				Any							    = num_stages
       _SCREAMING_SNAKE_CASE   :				int							    = hidden_sizes
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = depths
       _SCREAMING_SNAKE_CASE   :				Dict							    = is_training
       _SCREAMING_SNAKE_CASE   :				List[str]							    = use_labels
       _SCREAMING_SNAKE_CASE   :				int							    = intermediate_size
       _SCREAMING_SNAKE_CASE   :				Any							    = hidden_act
       _SCREAMING_SNAKE_CASE   :				Dict							    = type_sequence_label_size
       _SCREAMING_SNAKE_CASE   :				str							    = initializer_range
       _SCREAMING_SNAKE_CASE   :				List[str]							    = out_features
       _SCREAMING_SNAKE_CASE   :				Any							    = num_labels
       _SCREAMING_SNAKE_CASE   :				Dict							    = scope
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = num_stages
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				str							    = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = None
       if self.use_labels:
         _SCREAMING_SNAKE_CASE   :				List[str]							    = ids_tensor([self.batch_size]  ,     self.type_sequence_label_size				)

       _SCREAMING_SNAKE_CASE   :				Tuple							    = self.get_config()

       return config, pixel_values, labels
     def 				__SCREAMING_SNAKE_CASE							(     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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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




       """simple docstring"""


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




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = self.prepare_config_and_inputs()
       (
           (
           _SCREAMING_SNAKE_CASE
       )    ,		(
           _SCREAMING_SNAKE_CASE
       )    ,		(
           _SCREAMING_SNAKE_CASE
       )    ,		
       )   :				Optional[int]							    = config_and_inputs
       _SCREAMING_SNAKE_CASE   :				Tuple							    = {"pixel_values": pixel_values}
       return config, inputs_dict
@require_torch
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE      ,       unittest.TestCase  ):
     A__      =			(UperNetForSemanticSegmentation,) if is_torch_available() else ()
     A__      =			{"""image-segmentation""": UperNetForSemanticSegmentation} if is_torch_available() else {}
     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     A__      =			False
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = UperNetModelTester(self				)
       _SCREAMING_SNAKE_CASE   :				str							    = ConfigTester(self  ,     config_class=snake_case__  ,     has_text_modality=snake_case__  ,     hidden_size=37				)
     def 				__SCREAMING_SNAKE_CASE							(     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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       return
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.model_tester.prepare_config_and_inputs_for_common()

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

         _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = ["pixel_values"]
         self.assertListEqual(arg_names[:1]  ,     snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_semantic_segmentation(*snake_case__				)
     @unittest.skip(reason="UperNet does not use inputs_embeds"				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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




       """simple docstring"""


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




       """simple docstring"""


       pass
     @unittest.skip(reason="UperNet does not have a base model"				)
     def 				__SCREAMING_SNAKE_CASE							(     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 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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




       """simple docstring"""


       pass
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       def check_hidden_states_output(snake_case__  ,     snake_case__  ,     snake_case__				):
         _SCREAMING_SNAKE_CASE   :				Tuple							    = model_class(snake_case__				)
         model.to(snake_case__				)
         model.eval()

         with torch.no_grad():
           _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = model(**self._prepare_for_class(snake_case__  ,     snake_case__				)				)

         _SCREAMING_SNAKE_CASE   :				List[Any]							    = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

         _SCREAMING_SNAKE_CASE   :				Any							    = self.model_tester.num_stages
         self.assertEqual(len(snake_case__				)  ,     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    ,		_SCREAMING_SNAKE_CASE   :				Dict							    = self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         _SCREAMING_SNAKE_CASE   :				str							    = True
         check_hidden_states_output(snake_case__  ,     snake_case__  ,     snake_case__				)

         # check that output_hidden_states also work using config
         del inputs_dict["output_hidden_states"]
         _SCREAMING_SNAKE_CASE   :				List[str]							    = True

         check_hidden_states_output(snake_case__  ,     snake_case__  ,     snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				Dict							    = self.model_tester.prepare_config_and_inputs_for_common()

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




       """simple docstring"""


       pass




     @slow
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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





def    _lowerCAmelCase    (     )  ->    List[str]:
  _SCREAMING_SNAKE_CASE   :				Dict							    = hf_hub_download(
      repo_id="hf-internal-testing/fixtures_ade20k",						repo_type="dataset",						filename="ADE_val_00000001.jpg"				)
  _SCREAMING_SNAKE_CASE   :				Any							    = Image.open(lowerCamelCase__				).convert("RGB"				)
  return image





@require_torch
@require_vision
@slow
class 			UpperCamelCase      (							unittest.TestCase  ):
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny"				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny"				).to(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = prepare_img()
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = processor(images=snake_case__  ,     return_tensors="pt"				).to(snake_case__				)

       with torch.no_grad():
         _SCREAMING_SNAKE_CASE   :				Dict							    = model(**snake_case__				)

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = torch.Size((1, model.config.num_labels, 512, 512)				)
       self.assertEqual(outputs.logits.shape  ,     snake_case__				)

       _SCREAMING_SNAKE_CASE   :				Any							    = torch.tensor(
           [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]]				).to(snake_case__				)
       self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]  ,     snake_case__  ,     atol=1E-4				)				)




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny"				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny"				).to(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				Any							    = prepare_img()
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = processor(images=snake_case__  ,     return_tensors="pt"				).to(snake_case__				)

       with torch.no_grad():
         _SCREAMING_SNAKE_CASE   :				Optional[int]							    = model(**snake_case__				)

       _SCREAMING_SNAKE_CASE   :				int							    = torch.Size((1, model.config.num_labels, 512, 512)				)
       self.assertEqual(outputs.logits.shape  ,     snake_case__				)

       _SCREAMING_SNAKE_CASE   :				List[str]							    = torch.tensor(
           [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]				).to(snake_case__				)
       self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]  ,     snake_case__  ,     atol=1E-4				)				)





"""simple docstring"""
import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       unittest.TestCase  ):
     A__      =			DiTPipeline
     A__      =			CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
     A__      =			PipelineTesterMixin.required_optional_params - {
         """latents""",
         """num_images_per_prompt""",
         """callback""",
         """callback_steps""",
     }
     A__      =			CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
     A__      =			False
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       torch.manual_seed(0				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = TransformeraDModel(
           sample_size=16  ,     num_layers=2  ,     patch_size=4  ,     attention_head_dim=8  ,     num_attention_heads=2  ,     in_channels=4  ,     out_channels=8  ,     attention_bias=snake_case__  ,     activation_fn="gelu-approximate"  ,     num_embeds_ada_norm=1000  ,     norm_type="ada_norm_zero"  ,     norm_elementwise_affine=snake_case__  ,     )
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = AutoencoderKL()
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = DDIMScheduler()
       _SCREAMING_SNAKE_CASE   :				Dict							    = {"transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler}
       return components
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__=0				):




       """simple docstring"""


       if str(snake_case__				).startswith("mps"				):
         _SCREAMING_SNAKE_CASE   :				Tuple							    = torch.manual_seed(snake_case__				)
       else:
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = torch.Generator(device=snake_case__				).manual_seed(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = {
           "class_labels": [1],
           "generator": generator,
           "num_inference_steps": 2,
           "output_type": "numpy",
       }
       return inputs
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Dict							    = "cpu"

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.get_dummy_components()
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.pipeline_class(**snake_case__				)
       pipe.to(snake_case__				)
       pipe.set_progress_bar_config(disable=snake_case__				)

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = self.get_dummy_inputs(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = pipe(**snake_case__				).images
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = image[0, -3:, -3:, -1]

       self.assertEqual(image.shape  ,     (1, 16, 16, 3)				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = np.array([0.2_946, 0.6_601, 0.4_329, 0.3_296, 0.4_144, 0.5_319, 0.7_273, 0.5_013, 0.4_457]				)
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = np.abs(image_slice.flatten() - expected_slice				).max()
       self.assertLessEqual(snake_case__  ,     1E-3				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       self._test_inference_batch_single_identical(relax_max_difference=snake_case__  ,     expected_max_diff=1E-3				)




     @unittest.skipIf(
         torch_device != "cuda" or not is_xformers_available()  ,     reason="XFormers attention is only available with CUDA and `xformers` installed"  ,     )
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3				)
@require_torch_gpu
@slow
class 			UpperCamelCase      (							unittest.TestCase  ):
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				str							    = torch.manual_seed(0				)

       _SCREAMING_SNAKE_CASE   :				str							    = DiTPipeline.from_pretrained("facebook/DiT-XL-2-256"				)
       pipe.to("cuda"				)

       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = ["vase", "umbrella", "white shark", "white wolf"]
       _SCREAMING_SNAKE_CASE   :				Any							    = pipe.get_label_ids(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				List[str]							    = pipe(snake_case__  ,     generator=snake_case__  ,     num_inference_steps=40  ,     output_type="np"				).images

       for word, image in zip(snake_case__  ,     snake_case__				):
         _SCREAMING_SNAKE_CASE   :				int							    = load_numpy(
             F'''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy'''				)
         assert np.abs((expected_image - image).max()				) < 1E-2




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Any							    = DiTPipeline.from_pretrained("facebook/DiT-XL-2-512"				)
       _SCREAMING_SNAKE_CASE   :				Dict							    = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
       pipe.to("cuda"				)

       _SCREAMING_SNAKE_CASE   :				List[Any]							    = ["vase", "umbrella"]
       _SCREAMING_SNAKE_CASE   :				int							    = pipe.get_label_ids(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				str							    = torch.manual_seed(0				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = pipe(snake_case__  ,     generator=snake_case__  ,     num_inference_steps=25  ,     output_type="np"				).images

       for word, image in zip(snake_case__  ,     snake_case__				):
         _SCREAMING_SNAKE_CASE   :				Tuple							    = load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             F'''/dit/{word}_512.npy'''				)

         assert np.abs((expected_image - image).max()				) < 1E-1



