

from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import SwinConfig
from transformers.testing_utils import require_tf, require_vision, slow, to_atuple
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers.models.swin.modeling_tf_swin import (
           TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFSwinForImageClassification,
           TFSwinForMaskedImageModeling,
           TFSwinModel,
       )


if is_vision_available():
       from PIL import Image

       from transformers import AutoImageProcessor







class 		__magic_name__     :

   '''simple docstring'''




   def __init__(							self,       lowercase_,       lowercase_=13,       lowercase_=32,       lowercase_=2,       lowercase_=3,       lowercase_=16,       lowercase_=[1, 2, 1],       lowercase_=[2, 2, 4],       lowercase_=2,       lowercase_=2.0,       lowercase_=True,       lowercase_=0.0,       lowercase_=0.0,       lowercase_=0.1,       lowercase_="gelu",       lowercase_=False,       lowercase_=True,       lowercase_=0.02,       lowercase_=1E-5,       lowercase_=True,       lowercase_=None,       lowercase_=True,       lowercase_=10,       lowercase_=8,       )				->       None:






     """simple docstring"""






     a__    =parent
     a__    =batch_size
     a__    =image_size
     a__    =patch_size
     a__    =num_channels
     a__    =embed_dim
     a__    =depths
     a__    =num_heads
     a__    =window_size
     a__    =mlp_ratio
     a__    =qkv_bias
     a__    =hidden_dropout_prob
     a__    =attention_probs_dropout_prob
     a__    =drop_path_rate
     a__    =hidden_act
     a__    =use_absolute_embeddings
     a__    =patch_norm
     a__    =layer_norm_eps
     a__    =initializer_range
     a__    =is_training
     a__    =scope
     a__    =use_labels
     a__    =type_sequence_label_size
     a__    =encoder_stride



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     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.type_sequence_label_size			)

     a__    =self.get_config()

     return config, pixel_values, labels



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     return SwinConfig(
         image_size=self.image_size,       patch_size=self.patch_size,       num_channels=self.num_channels,       embed_dim=self.embed_dim,       depths=self.depths,       num_heads=self.num_heads,       window_size=self.window_size,       mlp_ratio=self.mlp_ratio,       qkv_bias=self.qkv_bias,       hidden_dropout_prob=self.hidden_dropout_prob,       attention_probs_dropout_prob=self.attention_probs_dropout_prob,       drop_path_rate=self.drop_path_rate,       hidden_act=self.hidden_act,       use_absolute_embeddings=self.use_absolute_embeddings,       path_norm=self.patch_norm,       layer_norm_eps=self.layer_norm_eps,       initializer_range=self.initializer_range,       encoder_stride=self.encoder_stride,       )



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       int:






     """simple docstring"""






     a__    =TFSwinModel(config=SCREAMING_SNAKE_CASE_			)
     a__    =model(SCREAMING_SNAKE_CASE_			)

     a__    =((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths			) - 1))
     a__    =int(config.embed_dim * 2 ** (len(config.depths			) - 1)			)

     self.parent.assertEqual(result.last_hidden_state.shape,       (self.batch_size, expected_seq_len, expected_dim)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       Any:






     """simple docstring"""






     a__    =TFSwinForMaskedImageModeling(config=SCREAMING_SNAKE_CASE_			)
     a__    =model(SCREAMING_SNAKE_CASE_			)
     self.parent.assertEqual(
         result.logits.shape,       (self.batch_size, self.num_channels, self.image_size, self.image_size)			)

     # test greyscale images
     a__    =1
     a__    =TFSwinForMaskedImageModeling(SCREAMING_SNAKE_CASE_			)

     a__    =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
     a__    =model(SCREAMING_SNAKE_CASE_			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, 1, self.image_size, self.image_size)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       Optional[int]:






     """simple docstring"""






     a__    =self.type_sequence_label_size
     a__    =TFSwinForImageClassification(SCREAMING_SNAKE_CASE_			)
     a__    =model(SCREAMING_SNAKE_CASE_,       labels=SCREAMING_SNAKE_CASE_			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.type_sequence_label_size)			)

     # test greyscale images
     a__    =1
     a__    =TFSwinForImageClassification(SCREAMING_SNAKE_CASE_			)
     a__    =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
     a__    =model(SCREAMING_SNAKE_CASE_			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.type_sequence_label_size)			)



   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






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







@require_tf
class 		__magic_name__     (				UpperCAmelCase_						,			UpperCAmelCase_						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      Union[str, Any]        =  (
       (
           TFSwinModel,
           TFSwinForImageClassification,
           TFSwinForMaskedImageModeling,
       )
       if is_tf_available()
       else ()
   )
   lowerCamelCase__					:      Tuple        =  (
       {"""feature-extraction""": TFSwinModel, """image-classification""": TFSwinForImageClassification}
       if is_tf_available()
       else {}
   )

   lowerCamelCase__					:      Any        =  False
   lowerCamelCase__					:      Tuple        =  False
   lowerCamelCase__					:      List[str]        =  False
   lowerCamelCase__					:      List[Any]        =  False



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__    =TFSwinModelTester(self			)
     a__    =ConfigTester(self,       config_class=SCREAMING_SNAKE_CASE_,       embed_dim=37			)



   def 					_UpperCAmelCase					(							self			)				->       str:






     """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 					_UpperCAmelCase					(							self			)				->       Optional[Any]:






     """simple docstring"""






     return



   def 					_UpperCAmelCase					(							self			)				->       Optional[int]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_			)



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_image_modeling(*SCREAMING_SNAKE_CASE_			)



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






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE_			)



   @unittest.skip(reason='''Swin does not use inputs_embeds'''			)
   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     pass



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       a__    =model_class(SCREAMING_SNAKE_CASE_			)
       self.assertIsInstance(model.get_input_embeddings(),       tf.keras.layers.Layer			)
       a__    =model.get_output_embeddings()
       self.assertTrue(x is None or isinstance(SCREAMING_SNAKE_CASE_,       tf.keras.layers.Dense			)			)



   def 					_UpperCAmelCase					(							self			)				->       Optional[int]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs_for_common()

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

       a__    =['''pixel_values''']
       self.assertListEqual(arg_names[:1],       SCREAMING_SNAKE_CASE_			)



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     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(SCREAMING_SNAKE_CASE_			)
       a__    =model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_			)			)
       a__    =outputs.attentions
       a__    =len(self.model_tester.depths			)
       self.assertEqual(len(SCREAMING_SNAKE_CASE_			),       SCREAMING_SNAKE_CASE_			)

       # check that output_attentions also work using config
       del inputs_dict["output_attentions"]
       a__    =True
       a__    =config.window_size**2
       a__    =model_class(SCREAMING_SNAKE_CASE_			)
       a__    =model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_			)			)
       a__    =outputs.attentions
       self.assertEqual(len(SCREAMING_SNAKE_CASE_			),       SCREAMING_SNAKE_CASE_			)

       self.assertListEqual(
           list(attentions[0].shape[-3:]			),       [self.model_tester.num_heads[0], window_size_squared, window_size_squared],       )
       a__    =len(SCREAMING_SNAKE_CASE_			)

       # Check attention is always last and order is fine
       a__    =True
       a__    =True
       a__    =model_class(SCREAMING_SNAKE_CASE_			)
       a__    =model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_			)			)

       if hasattr(self.model_tester,       '''num_hidden_states_types'''			):
         a__    =self.model_tester.num_hidden_states_types
       else:
         # also another +1 for reshaped_hidden_states
         a__    =2
       self.assertEqual(out_len + added_hidden_states,       len(SCREAMING_SNAKE_CASE_			)			)

       a__    =outputs.attentions

       self.assertEqual(len(SCREAMING_SNAKE_CASE_			),       SCREAMING_SNAKE_CASE_			)

       self.assertListEqual(
           list(self_attentions[0].shape[-3:]			),       [self.model_tester.num_heads[0], window_size_squared, window_size_squared],       )



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_			)				->       str:






     """simple docstring"""






     a__    =model_class(SCREAMING_SNAKE_CASE_			)
     a__    =model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_			)			)
     a__    =outputs.hidden_states

     a__    =getattr(
         self.model_tester,       '''expected_num_hidden_layers''',       len(self.model_tester.depths			) + 1			)
     self.assertEqual(len(SCREAMING_SNAKE_CASE_			),       SCREAMING_SNAKE_CASE_			)

     # Swin has a different seq_length
     a__    =to_atuple(config.patch_size			)

     a__    =(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

     self.assertListEqual(
         list(hidden_states[0].shape[-2:]			),       [num_patches, self.model_tester.embed_dim],       )

     a__    =outputs.reshaped_hidden_states
     self.assertEqual(len(SCREAMING_SNAKE_CASE_			),       SCREAMING_SNAKE_CASE_			)

     a__    =reshaped_hidden_states[0].shape

     a__    =tf.reshape(reshaped_hidden_states[0],       (batch_size, num_channels, height * width)			)
     a__    =tf.transpose(SCREAMING_SNAKE_CASE_,       (0, 2, 1)			)

     self.assertListEqual(
         list(reshaped_hidden_states.shape[-2:]			),       [num_patches, self.model_tester.embed_dim],       )



   def 					_UpperCAmelCase					(							self			)				->       Optional[int]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs_for_common()

     a__    =to_atuple(self.model_tester.image_size			)

     for model_class in self.all_model_classes:
       a__    =True
       self.check_hidden_states_output(SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_			)

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

       self.check_hidden_states_output(SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_			)



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs_for_common()
     a__    =3

     a__    =to_atuple(self.model_tester.image_size			)
     a__    =to_atuple(config.patch_size			)

     a__    =image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
     a__    =image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

     for model_class in self.all_model_classes:
       a__    =True
       self.check_hidden_states_output(SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_,       (padded_height, padded_width)			)

       # check that output_hidden_states also work using config
       del inputs_dict["output_hidden_states"]
       a__    =True
       self.check_hidden_states_output(SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_,       SCREAMING_SNAKE_CASE_,       (padded_height, padded_width)			)



   @slow
   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     for model_name in TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       a__    =TFSwinModel.from_pretrained(SCREAMING_SNAKE_CASE_			)
       self.assertIsNotNone(SCREAMING_SNAKE_CASE_			)







@require_vision
@require_tf
class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




   @cached_property
   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     return (
         AutoImageProcessor.from_pretrained('''microsoft/swin-tiny-patch4-window7-224'''			)
         if is_vision_available()
         else None
     )



   @slow
   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =TFSwinForImageClassification.from_pretrained('''microsoft/swin-tiny-patch4-window7-224'''			)
     a__    =self.default_image_processor

     a__    =Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''			)
     a__    =image_processor(images=SCREAMING_SNAKE_CASE_,       return_tensors='''tf'''			)

     # forward pass
     a__    =model(SCREAMING_SNAKE_CASE_			)

     # verify the logits
     a__    =tf.TensorShape((1, 1000)			)
     self.assertEqual(outputs.logits.shape,       SCREAMING_SNAKE_CASE_			)
     a__    =tf.constant([-0.0948, -0.6454, -0.0921]			)
     self.assertTrue(np.allclose(outputs.logits[0, :3],       SCREAMING_SNAKE_CASE_,       atol=1E-4			)			)






'''simple docstring'''


import gc
import unittest

import numpy as np
import torch

from diffusers import RePaintPipeline, RePaintScheduler, UNetaDModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_image,
    load_numpy,
    nightly,
    require_torch_gpu,
    skip_mps,
    torch_device,
)

from ..pipeline_params import IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class 						lowerCAmelCase_	(       UpperCAmelCase_       ,					unittest.TestCase							):







  '''simple docstring'''

  UpperCamelCase_			:      Any										=				RePaintPipeline
  UpperCamelCase_			:      int										=				IMAGE_INPAINTING_PARAMS - {"""width""", """height""", """guidance_scale"""}
  UpperCamelCase_			:      List[str]										=				PipelineTesterMixin.required_optional_params - {
      """latents""",
      """num_images_per_prompt""",
      """callback""",
      """callback_steps""",
  }
  UpperCamelCase_			:      List[str]										=				IMAGE_INPAINTING_BATCH_PARAMS
  UpperCamelCase_			:      str										=				False



  def 					_snake_case     ( self				:							Dict	)			->		Optional[Any]:







       '''simple docstring'''



       torch.manual_seed(0	)
       torch.manual_seed(0	)
       A:		Optional[Any]          =		UNetaDModel(
           block_out_channels=(32, 64)					,  layers_per_block=2					,  sample_size=32					,  in_channels=3					,  out_channels=3					,  down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')					,  up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')					,  )
       A:		Optional[int]          =		RePaintScheduler()
       A:		Optional[int]          =		{'''unet''': unet, '''scheduler''': scheduler}
       return components



  def 					_snake_case     ( self				:							Optional[Any]					,  SCREAMING_SNAKE_CASE_				:							List[Any]					,  SCREAMING_SNAKE_CASE_				:							Tuple=0	)			->		Optional[int]:







       '''simple docstring'''



       if str(SCREAMING_SNAKE_CASE_	).startswith('''mps'''	):
            A:		str          =		torch.manual_seed(SCREAMING_SNAKE_CASE_	)
       else:
            A:		Any          =		torch.Generator(device=SCREAMING_SNAKE_CASE_	).manual_seed(SCREAMING_SNAKE_CASE_	)
       A:		str          =		np.random.RandomState(SCREAMING_SNAKE_CASE_	).standard_normal((1, 3, 32, 32)	)
       A:		Tuple          =		torch.from_numpy(SCREAMING_SNAKE_CASE_	).to(device=SCREAMING_SNAKE_CASE_					,  dtype=torch.floataa	)
       A:		List[Any]          =		(image > 0).to(device=SCREAMING_SNAKE_CASE_					,  dtype=torch.floataa	)
       A:		Tuple          =		{
           '''image''': image,
           '''mask_image''': mask,
           '''generator''': generator,
           '''num_inference_steps''': 5,
           '''eta''': 0.0,
           '''jump_length''': 2,
           '''jump_n_sample''': 2,
           '''output_type''': '''numpy''',
       }
       return inputs



  def 					_snake_case     ( self				:							int	)			->		Tuple:







       '''simple docstring'''



       A:		int          =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
       A:		List[str]          =		self.get_dummy_components()
       A:		int          =		RePaintPipeline(**SCREAMING_SNAKE_CASE_	)
       A:		List[Any]          =		sd_pipe.to(SCREAMING_SNAKE_CASE_	)
       sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)

       A:		Dict          =		self.get_dummy_inputs(SCREAMING_SNAKE_CASE_	)
       A:		List[str]          =		sd_pipe(**SCREAMING_SNAKE_CASE_	).images
       A:		Optional[Any]          =		image[0, -3:, -3:, -1]

       assert image.shape == (1, 32, 32, 3)
       A:		List[Any]          =		np.array([1.0000, 0.5426, 0.5497, 0.2200, 1.0000, 1.0000, 0.5623, 1.0000, 0.6274]	)

       assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-3



  @skip_mps
  def 					_snake_case     ( self				:							Union[str, Any]	)			->		Optional[Any]:







       '''simple docstring'''



       return super().test_save_load_local()



  @unittest.skip('''non-deterministic pipeline'''	)
  def 					_snake_case     ( self				:							Tuple	)			->		Optional[Any]:







       '''simple docstring'''



       return super().test_inference_batch_single_identical()



  @skip_mps
  def 					_snake_case     ( self				:							List[Any]	)			->		str:







       '''simple docstring'''



       return super().test_dict_tuple_outputs_equivalent()



  @skip_mps
  def 					_snake_case     ( self				:							List[Any]	)			->		str:







       '''simple docstring'''



       return super().test_save_load_optional_components()




  @skip_mps
  def 					_snake_case     ( self				:							Any	)			->		Optional[Any]:







       '''simple docstring'''



       return super().test_attention_slicing_forward_pass()






@nightly
@require_torch_gpu
class 						lowerCAmelCase_	(       unittest.TestCase							):







  '''simple docstring'''

  def 					_snake_case     ( self				:							Optional[Any]	)			->		Optional[Any]:







       '''simple docstring'''



       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




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







       '''simple docstring'''



       A:		str          =		load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'''
           '''repaint/celeba_hq_256.png'''	)
       A:		Union[str, Any]          =		load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png'''	)
       A:		Optional[Any]          =		load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'''
           '''repaint/celeba_hq_256_result.npy'''	)

       A:		Dict          =		'''google/ddpm-ema-celebahq-256'''
       A:		Tuple          =		UNetaDModel.from_pretrained(SCREAMING_SNAKE_CASE_	)
       A:		List[Any]          =		RePaintScheduler.from_pretrained(SCREAMING_SNAKE_CASE_	)

       A:		Optional[Any]          =		RePaintPipeline(unet=SCREAMING_SNAKE_CASE_					,  scheduler=SCREAMING_SNAKE_CASE_	).to(SCREAMING_SNAKE_CASE_	)
       repaint.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)
       repaint.enable_attention_slicing()

       A:		str          =		torch.manual_seed(0	)
       A:		str          =		repaint(
           SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  num_inference_steps=2_50					,  eta=0.0					,  jump_length=10					,  jump_n_sample=10					,  generator=SCREAMING_SNAKE_CASE_					,  output_type='''np'''					,  )
       A:		Optional[Any]          =		output.images[0]

       assert image.shape == (2_56, 2_56, 3)
       assert np.abs(expected_image - image	).mean() < 1E-2






