


import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           MODEL_MAPPING,
           SegformerForImageClassification,
           SegformerForSemanticSegmentation,
           SegformerModel,
       )
       from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image

       from transformers import SegformerImageProcessor






class 				__A    (	UpperCamelCase__      ):






 """simple docstring"""




 def 				SCREAMING_SNAKE_CASE							(		self						) ->					str:
     a					=self.config_class(**self.inputs_dict						)
     self.parent.assertTrue(hasattr(__lowerCamelCase			,   '''hidden_sizes'''						)						)
     self.parent.assertTrue(hasattr(__lowerCamelCase			,   '''num_attention_heads'''						)						)
     self.parent.assertTrue(hasattr(__lowerCamelCase			,   '''num_encoder_blocks'''						)						)






class 				__A    :






 """simple docstring"""




 def __init__(		self			,   __A			,   __A=13			,   __A=64			,   __A=3			,   __A=4			,   __A=[2, 2, 2, 2]			,   __A=[8, 4, 2, 1]			,   __A=[16, 32, 64, 128]			,   __A=[1, 4, 8, 16]			,   __A=[1, 2, 4, 8]			,   __A=True			,   __A=True			,   __A="gelu"			,   __A=0.1			,   __A=0.1			,   __A=0.02			,   __A=3			,   __A=None			,   ) ->					Dict:
     a					=parent
     a					=batch_size
     a					=image_size
     a					=num_channels
     a					=num_encoder_blocks
     a					=sr_ratios
     a					=depths
     a					=hidden_sizes
     a					=downsampling_rates
     a					=num_attention_heads
     a					=is_training
     a					=use_labels
     a					=hidden_act
     a					=hidden_dropout_prob
     a					=attention_probs_dropout_prob
     a					=initializer_range
     a					=num_labels
     a					=scope


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Any:
     a					=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

     a					=None
     if self.use_labels:
         a					=ids_tensor([self.batch_size, self.image_size, self.image_size]			,   self.num_labels						)

     a					=self.get_config()
     return config, pixel_values, labels


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Union[str, Any]:
     return SegformerConfig(
         image_size=self.image_size			,   num_channels=self.num_channels			,   num_encoder_blocks=self.num_encoder_blocks			,   depths=self.depths			,   hidden_sizes=self.hidden_sizes			,   num_attention_heads=self.num_attention_heads			,   hidden_act=self.hidden_act			,   hidden_dropout_prob=self.hidden_dropout_prob			,   attention_probs_dropout_prob=self.attention_probs_dropout_prob			,   initializer_range=self.initializer_range			,   )


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A						) ->					Optional[int]:
     a					=SegformerModel(config=__lowerCamelCase						)
     model.to(__lowerCamelCase						)
     model.eval()
     a					=model(__lowerCamelCase						)
     a					=a					=self.image_size // (self.downsampling_rates[-1] * 2)
     self.parent.assertEqual(
         result.last_hidden_state.shape			,   (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)						)


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A						) ->					str:
     a					=self.num_labels
     a					=SegformerForSemanticSegmentation(__lowerCamelCase						)
     model.to(__lowerCamelCase						)
     model.eval()
     a					=model(__lowerCamelCase						)
     self.parent.assertEqual(
         result.logits.shape			,   (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)						)
     a					=model(__lowerCamelCase			,   labels=__lowerCamelCase						)
     self.parent.assertEqual(
         result.logits.shape			,   (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)						)
     self.parent.assertGreater(result.loss			,   0.0						)


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A			,   __A						) ->					Dict:
     a					=1
     a					=SegformerForSemanticSegmentation(config=__lowerCamelCase						)
     model.to(__lowerCamelCase						)
     model.eval()
     a					=torch.randint(0			,   1			,   (self.batch_size, self.image_size, self.image_size)						).to(__lowerCamelCase						)
     a					=model(__lowerCamelCase			,   labels=__lowerCamelCase						)
     self.parent.assertGreater(result.loss			,   0.0						)





 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
     a					=self.prepare_config_and_inputs()
     a		,				a		,				a					=config_and_inputs
     a					={'''pixel_values''': pixel_values}
     return config, inputs_dict






@require_torch
class 				__A    (	UpperCamelCase__,   UpperCamelCase__,   unittest.TestCase      ):






 """simple docstring"""




 __lowerCAmelCase			 =  (
     (
         SegformerModel,
         SegformerForSemanticSegmentation,
         SegformerForImageClassification,
     )
     if is_torch_available()
     else ()
 )
 __lowerCAmelCase			 =  (
     {
         "feature-extraction": SegformerModel,
         "image-classification": SegformerForImageClassification,
         "image-segmentation": SegformerForSemanticSegmentation,
     }
     if is_torch_available()
     else {}
 )

 __lowerCAmelCase			 =  True
 __lowerCAmelCase			 =  False
 __lowerCAmelCase			 =  False
 __lowerCAmelCase			 =  False


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[Any]:
     a					=SegformerModelTester(self						)
     a					=SegformerConfigTester(self			,   config_class=__lowerCamelCase						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
     self.config_tester.run_common_tests()


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Tuple:
     a					=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__lowerCamelCase						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Any:
     a					=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_binary_image_segmentation(*__lowerCamelCase						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Union[str, Any]:
     a					=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_segmentation(*__lowerCamelCase						)


 @unittest.skip('''SegFormer does not use inputs_embeds'''						)
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Any:
     pass


 @unittest.skip('''SegFormer does not have get_input_embeddings method and get_output_embeddings methods'''						)
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Any:
     pass


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[str]:
     a		,				a					=self.model_tester.prepare_config_and_inputs_for_common()

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

         a					=['''pixel_values''']
         self.assertListEqual(arg_names[:1]			,   __lowerCamelCase						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					int:
     a		,				a					=self.model_tester.prepare_config_and_inputs_for_common()
     a					=True

     for model_class in self.all_model_classes:
         a					=True
         a					=False
         a					=True
         a					=model_class(__lowerCamelCase						)
         model.to(__lowerCamelCase						)
         model.eval()
         with torch.no_grad():
             a					=model(**self._prepare_for_class(__lowerCamelCase			,   __lowerCamelCase						)						)
         a					=outputs.attentions

         a					=sum(self.model_tester.depths						)
         self.assertEqual(len(__lowerCamelCase						)			,   __lowerCamelCase						)

         # check that output_attentions also work using config
         del inputs_dict["output_attentions"]
         a					=True
         a					=model_class(__lowerCamelCase						)
         model.to(__lowerCamelCase						)
         model.eval()
         with torch.no_grad():
             a					=model(**self._prepare_for_class(__lowerCamelCase			,   __lowerCamelCase						)						)
         a					=outputs.attentions

         self.assertEqual(len(__lowerCamelCase						)			,   __lowerCamelCase						)

         # verify the first attentions (first block, first layer)
         a					=(self.model_tester.image_size // 4) ** 2
         a					=(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
         self.assertListEqual(
             list(attentions[0].shape[-3:]						)			,   [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]			,   )

         # verify the last attentions (last block, last layer)
         a					=(self.model_tester.image_size // 32) ** 2
         a					=(self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
         self.assertListEqual(
             list(attentions[-1].shape[-3:]						)			,   [self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]			,   )
         a					=len(__lowerCamelCase						)

         # Check attention is always last and order is fine
         a					=True
         a					=True
         a					=model_class(__lowerCamelCase						)
         model.to(__lowerCamelCase						)
         model.eval()
         with torch.no_grad():
             a					=model(**self._prepare_for_class(__lowerCamelCase			,   __lowerCamelCase						)						)

         self.assertEqual(out_len + 1			,   len(__lowerCamelCase						)						)

         a					=outputs.attentions

         self.assertEqual(len(__lowerCamelCase						)			,   __lowerCamelCase						)
         # verify the first attentions (first block, first layer)
         a					=(self.model_tester.image_size // 4) ** 2
         a					=(self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
         self.assertListEqual(
             list(self_attentions[0].shape[-3:]						)			,   [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]			,   )


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Union[str, Any]:
     def check_hidden_states_output(__A			,   __A			,   __A						):
         a					=model_class(__lowerCamelCase						)
         model.to(__lowerCamelCase						)
         model.eval()

         with torch.no_grad():
             a					=model(**self._prepare_for_class(__lowerCamelCase			,   __lowerCamelCase						)						)

         a					=outputs.hidden_states

         a					=self.model_tester.num_encoder_blocks
         self.assertEqual(len(__lowerCamelCase						)			,   __lowerCamelCase						)

         # verify the first hidden states (first block)
         self.assertListEqual(
             list(hidden_states[0].shape[-3:]						)			,   [
                 self.model_tester.hidden_sizes[0],
                 self.model_tester.image_size // 4,
                 self.model_tester.image_size // 4,
             ]			,   )

     a		,				a					=self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
         a					=True
         check_hidden_states_output(__lowerCamelCase			,   __lowerCamelCase			,   __lowerCamelCase						)

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

         check_hidden_states_output(__lowerCamelCase			,   __lowerCamelCase			,   __lowerCamelCase						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[Any]:
     if not self.model_tester.is_training:
         return

     a		,				a					=self.model_tester.prepare_config_and_inputs_for_common()
     a					=True

     for model_class in self.all_model_classes:
         if model_class in get_values(__lowerCamelCase						):
             continue

         a					=model_class(__lowerCamelCase						)
         model.to(__lowerCamelCase						)
         model.train()
         a					=self._prepare_for_class(__lowerCamelCase			,   __lowerCamelCase			,   return_labels=__lowerCamelCase						)
         a					=model(**__lowerCamelCase						).loss
         loss.backward()


 @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''						)
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Union[str, Any]:
     pass





 @slow
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
     for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         a					=SegformerModel.from_pretrained(__lowerCamelCase						)
         self.assertIsNotNone(__lowerCamelCase						)

def _A						(						):







    """simple docstring"""
    a					=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''		)
    return image






@require_torch
class 				__A    (	unittest.TestCase      ):






 """simple docstring"""




 @slow
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[Any]:
     # only resize + normalize
     a					=SegformerImageProcessor(
         image_scale=(512, 512)			,   keep_ratio=__lowerCamelCase			,   align=__lowerCamelCase			,   do_random_crop=__lowerCamelCase						)
     a					=SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''						).to(
         __lowerCamelCase						)

     a					=prepare_img()
     a					=image_processor(images=__lowerCamelCase			,   return_tensors='''pt'''						)
     a					=encoded_inputs.pixel_values.to(__lowerCamelCase						)

     with torch.no_grad():
         a					=model(__lowerCamelCase						)

     a					=torch.Size((1, model.config.num_labels, 128, 128)						)
     self.assertEqual(outputs.logits.shape			,   __lowerCamelCase						)

     a					=torch.tensor(
         [
             [[-4.6_310, -5.5_232, -6.2_356], [-5.1_921, -6.1_444, -6.5_996], [-5.4_424, -6.2_790, -6.7_574]],
             [[-12.1_391, -13.3_122, -13.9_554], [-12.8_732, -13.9_352, -14.3_563], [-12.9_438, -13.8_226, -14.2_513]],
             [[-12.5_134, -13.4_686, -14.4_915], [-12.8_669, -14.4_343, -14.7_758], [-13.2_523, -14.5_819, -15.0_694]],
         ]						).to(__lowerCamelCase						)
     self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]			,   __lowerCamelCase			,   atol=1E-4						)						)


 @slow
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[str]:
     # only resize + normalize
     a					=SegformerImageProcessor(
         image_scale=(512, 512)			,   keep_ratio=__lowerCamelCase			,   align=__lowerCamelCase			,   do_random_crop=__lowerCamelCase						)
     a					=SegformerForSemanticSegmentation.from_pretrained(
         '''nvidia/segformer-b1-finetuned-cityscapes-1024-1024'''						).to(__lowerCamelCase						)

     a					=prepare_img()
     a					=image_processor(images=__lowerCamelCase			,   return_tensors='''pt'''						)
     a					=encoded_inputs.pixel_values.to(__lowerCamelCase						)

     with torch.no_grad():
         a					=model(__lowerCamelCase						)

     a					=torch.Size((1, model.config.num_labels, 128, 128)						)
     self.assertEqual(outputs.logits.shape			,   __lowerCamelCase						)

     a					=torch.tensor(
         [
             [[-13.5_748, -13.9_111, -12.6_500], [-14.3_500, -15.3_683, -14.2_328], [-14.7_532, -16.0_424, -15.6_087]],
             [[-17.1_651, -15.8_725, -12.9_653], [-17.2_580, -17.3_718, -14.8_223], [-16.6_058, -16.8_783, -16.7_452]],
             [[-3.6_456, -3.0_209, -1.4_203], [-3.0_797, -3.1_959, -2.0_000], [-1.8_757, -1.9_217, -1.6_997]],
         ]						).to(__lowerCamelCase						)
     self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]			,   __lowerCamelCase			,   atol=1E-1						)						)





 @slow
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[str]:
     # only resize + normalize
     a					=SegformerImageProcessor(
         image_scale=(512, 512)			,   keep_ratio=__lowerCamelCase			,   align=__lowerCamelCase			,   do_random_crop=__lowerCamelCase						)
     a					=SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''						).to(
         __lowerCamelCase						)

     a					=prepare_img()
     a					=image_processor(images=__lowerCamelCase			,   return_tensors='''pt'''						)
     a					=encoded_inputs.pixel_values.to(__lowerCamelCase						)

     with torch.no_grad():
         a					=model(__lowerCamelCase						)

     a					=outputs.logits.detach().cpu()

     a					=image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase			,   target_sizes=[(500, 300)]						)
     a					=torch.Size((500, 300)						)
     self.assertEqual(segmentation[0].shape			,   __lowerCamelCase						)

     a					=image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase						)
     a					=torch.Size((128, 128)						)
     self.assertEqual(segmentation[0].shape			,   __lowerCamelCase						)



"""simple docstring"""







import random
import unittest

import torch

from diffusers import IFInpaintingSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin






@skip_mps
class 				__A    (	_SCREAMING_SNAKE_CASE,   _SCREAMING_SNAKE_CASE,   unittest.TestCase      ):






 """simple docstring"""




 __lowerCAmelCase			 =  IFInpaintingSuperResolutionPipeline
 __lowerCAmelCase			 =  TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"width", "height"}
 __lowerCAmelCase			 =  TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS.union({"original_image"}      )
 __lowerCAmelCase			 =  PipelineTesterMixin.required_optional_params - {"latents"}


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Optional[int]:
     return self._get_superresolution_dummy_components()


 def 				SCREAMING_SNAKE_CASE							(		self			,   __A			,   __A=0						) ->					Optional[int]:
     if str(__A						).startswith('''mps'''						):
         a					=torch.manual_seed(__A						)
     else:
         a					=torch.Generator(device=__A						).manual_seed(__A						)

     a					=floats_tensor((1, 3, 16, 16)			,   rng=random.Random(__A						)						).to(__A						)
     a					=floats_tensor((1, 3, 32, 32)			,   rng=random.Random(__A						)						).to(__A						)
     a					=floats_tensor((1, 3, 32, 32)			,   rng=random.Random(__A						)						).to(__A						)

     a					={
         '''prompt''': '''A painting of a squirrel eating a burger''',
         '''image''': image,
         '''original_image''': original_image,
         '''mask_image''': mask_image,
         '''generator''': generator,
         '''num_inference_steps''': 2,
         '''output_type''': '''numpy''',
     }

     return inputs


 @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						) ->					List[Any]:
     self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Dict:
     self._test_save_load_optional_components()


 @unittest.skipIf(torch_device != '''cuda'''			,   reason='''float16 requires CUDA'''						)
 def 				SCREAMING_SNAKE_CASE							(		self						) ->					int:
     # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
     super().test_save_load_floataa(expected_max_diff=1E-1						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Any:
     self._test_attention_slicing_forward_pass(expected_max_diff=1E-2						)


 def 				SCREAMING_SNAKE_CASE							(		self						) ->					List[Any]:
     self._test_save_load_local()





 def 				SCREAMING_SNAKE_CASE							(		self						) ->					Dict:
     self._test_inference_batch_single_identical(
         expected_max_diff=1E-2			,   )