






"""simple docstring"""






import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
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 torch import nn

  from transformers import (
      FocalNetBackbone,
      FocalNetForImageClassification,
      FocalNetForMaskedImageModeling,
      FocalNetModel,
  )
  from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor




class 							lowerCAmelCase__				:






      '''simple docstring'''

      def __init__(						self     :  str						,	lowercase_     :  str						,	lowercase_     :  Union[str, Any]=13						,	lowercase_     :  Dict=32						,	lowercase_     :  str=2						,	lowercase_     :  Union[str, Any]=3						,	lowercase_     :  List[str]=16						,	lowercase_     :  int=[32, 64, 128]						,	lowercase_     :  List[Any]=[1, 2, 1]						,	lowercase_     :  int=[2, 2, 4]						,	lowercase_     :  str=2						,	lowercase_     :  Optional[Any]=2.0						,	lowercase_     :  List[str]=True						,	lowercase_     :  List[str]=0.0						,	lowercase_     :  List[Any]=0.0						,	lowercase_     :  Tuple=0.1						,	lowercase_     :  Union[str, Any]="gelu"						,	lowercase_     :  int=False						,	lowercase_     :  Dict=True						,	lowercase_     :  int=0.02						,	lowercase_     :  Optional[Any]=1e-5						,	lowercase_     :  int=True						,	lowercase_     :  Union[str, Any]=None						,	lowercase_     :  List[str]=True						,	lowercase_     :  Tuple=10						,	lowercase_     :  Any=8						,	lowercase_     :  Any=["stage1", "stage2"]						,	lowercase_     :  str=[1, 2]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Any    						=		parent
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		batch_size
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		image_size
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		patch_size
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		num_channels
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		embed_dim
          SCREAMING_SNAKE_CASE_      :  Dict    						=		hidden_sizes
          SCREAMING_SNAKE_CASE_      :  str    						=		depths
          SCREAMING_SNAKE_CASE_      :  Any    						=		num_heads
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		window_size
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		mlp_ratio
          SCREAMING_SNAKE_CASE_      :  Any    						=		qkv_bias
          SCREAMING_SNAKE_CASE_      :  Dict    						=		hidden_dropout_prob
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		attention_probs_dropout_prob
          SCREAMING_SNAKE_CASE_      :  str    						=		drop_path_rate
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		hidden_act
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		use_absolute_embeddings
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		patch_norm
          SCREAMING_SNAKE_CASE_      :  Any    						=		layer_norm_eps
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		initializer_range
          SCREAMING_SNAKE_CASE_      :  str    						=		is_training
          SCREAMING_SNAKE_CASE_      :  Any    						=		scope
          SCREAMING_SNAKE_CASE_      :  str    						=		use_labels
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		type_sequence_label_size
          SCREAMING_SNAKE_CASE_      :  str    						=		encoder_stride
          SCREAMING_SNAKE_CASE_      :  Any    						=		out_features
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		out_indices
      def _SCREAMING_SNAKE_CASE (						self     :  Any):




          '''simple docstring'''


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

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

          SCREAMING_SNAKE_CASE_      :  str    						=		self.get_config()

          return config, pixel_values, labels
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          return FocalNetConfig(
              image_size=self.image_size						,	patch_size=self.patch_size						,	num_channels=self.num_channels						,	embed_dim=self.embed_dim						,	hidden_sizes=self.hidden_sizes						,	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						,	out_features=self.out_features						,	out_indices=self.out_indices						,	)
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple						,	lowercase_     :  List[str]						,	lowercase_     :  Dict						,	lowercase_     :  Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		FocalNetModel(config=lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		model(lowercase_)

          SCREAMING_SNAKE_CASE_      :  Tuple    						=		((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths) - 1))
          SCREAMING_SNAKE_CASE_      :  Any    						=		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 _SCREAMING_SNAKE_CASE (						self     :  int						,	lowercase_     :  Optional[Any]						,	lowercase_     :  List[Any]						,	lowercase_     :  Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		FocalNetBackbone(config=lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  int    						=		model(lowercase_)

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps)						,	len(config.out_features))
          self.parent.assertListEqual(list(result.feature_maps[0].shape)						,	[self.batch_size, self.image_size, 8, 8])

          # verify channels
          self.parent.assertEqual(len(model.channels)						,	len(config.out_features))
          self.parent.assertListEqual(model.channels						,	config.hidden_sizes[:-1])

          # verify backbone works with out_features=None
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		None
          SCREAMING_SNAKE_CASE_      :  str    						=		FocalNetBackbone(config=lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  str    						=		model(lowercase_)

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps)						,	1)
          self.parent.assertListEqual(list(result.feature_maps[0].shape)						,	[self.batch_size, self.image_size * 2, 4, 4])

          # verify channels
          self.parent.assertEqual(len(model.channels)						,	1)
          self.parent.assertListEqual(model.channels						,	[config.hidden_sizes[-1]])
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]						,	lowercase_     :  Union[str, Any]						,	lowercase_     :  List[Any]						,	lowercase_     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		FocalNetForMaskedImageModeling(config=lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		model(lowercase_)
          self.parent.assertEqual(
              result.reconstruction.shape						,	(self.batch_size, self.num_channels, self.image_size, self.image_size))

          # test greyscale images
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		1
          SCREAMING_SNAKE_CASE_      :  Any    						=		FocalNetForMaskedImageModeling(lowercase_)
          model.to(lowercase_)
          model.eval()

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
          SCREAMING_SNAKE_CASE_      :  Dict    						=		model(lowercase_)
          self.parent.assertEqual(result.reconstruction.shape						,	(self.batch_size, 1, self.image_size, self.image_size))
      def _SCREAMING_SNAKE_CASE (						self     :  int						,	lowercase_     :  Tuple						,	lowercase_     :  str						,	lowercase_     :  str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  int    						=		self.type_sequence_label_size
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		FocalNetForImageClassification(lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		model(lowercase_						,	labels=lowercase_)
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.type_sequence_label_size))

          # test greyscale images
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		1
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		FocalNetForImageClassification(lowercase_)
          model.to(lowercase_)
          model.eval()

          SCREAMING_SNAKE_CASE_      :  int    						=		floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
          SCREAMING_SNAKE_CASE_      :  str    						=		model(lowercase_)
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.type_sequence_label_size))



      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  int    						=		self.prepare_config_and_inputs()

          SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  List[str]    						=		config_and_inputs
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		{'''pixel_values''': pixel_values}
          return config, inputs_dict




@require_torch
class 							lowerCAmelCase__				(			UpperCAmelCase__				,		UpperCAmelCase__				,		unittest.TestCase ):






      '''simple docstring'''

      __UpperCamelCase	 =  (
          (
              FocalNetModel,
              FocalNetForImageClassification,
              FocalNetForMaskedImageModeling,
              FocalNetBackbone,
          )
          if is_torch_available()
          else ()
      )
      __UpperCamelCase	 =  (
          {"feature-extraction": FocalNetModel, "image-classification": FocalNetForImageClassification}
          if is_torch_available()
          else {}
      )
      __UpperCamelCase	 =  False

      __UpperCamelCase	 =  False
      __UpperCamelCase	 =  False
      __UpperCamelCase	 =  False
      __UpperCamelCase	 =  False
      def _SCREAMING_SNAKE_CASE (						self     :  Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		FocalNetModelTester(self)
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		ConfigTester(self						,	config_class=lowercase_						,	embed_dim=37						,	has_text_modality=lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]):




          '''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     :  List[str]):




          '''simple docstring'''


          return
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_backbone(*lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_image_modeling(*lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  str    						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*lowercase_)
      @unittest.skip(reason='''FocalNet does not use inputs_embeds''')
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]):




          '''simple docstring'''


          pass
      @unittest.skip(reason='''FocalNet does not use feedforward chunking''')
      def _SCREAMING_SNAKE_CASE (						self     :  str):




          '''simple docstring'''


          pass
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''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[:-1]:
              SCREAMING_SNAKE_CASE_      :  Any    						=		model_class(lowercase_)
              self.assertIsInstance(model.get_input_embeddings()						,	(nn.Module))
              SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		model.get_output_embeddings()
              self.assertTrue(x is None or isinstance(lowercase_						,	nn.Linear))
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  List[str]    						=		self.model_tester.prepare_config_and_inputs_for_common()

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

              SCREAMING_SNAKE_CASE_      :  List[Any]    						=		['''pixel_values''']
              self.assertListEqual(arg_names[:1]						,	lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  str						,	lowercase_     :  str						,	lowercase_     :  List[str]						,	lowercase_     :  Optional[int]						,	lowercase_     :  List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		model_class(lowercase_)
          model.to(lowercase_)
          model.eval()

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

          SCREAMING_SNAKE_CASE_      :  List[str]    						=		outputs.hidden_states

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		getattr(
              self.model_tester						,	'''expected_num_hidden_layers'''						,	len(self.model_tester.depths) + 1)
          self.assertEqual(len(lowercase_)						,	lowercase_)

          # FocalNet has a different seq_length
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		(
              config.patch_size
              if isinstance(config.patch_size						,	collections.abc.Iterable)
              else (config.patch_size, config.patch_size)
          )

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		(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]						,	)

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		outputs.reshaped_hidden_states
          self.assertEqual(len(lowercase_)						,	lowercase_)

          SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  List[Any]    						=		reshaped_hidden_states[0].shape
          SCREAMING_SNAKE_CASE_      :  int    						=		(
              reshaped_hidden_states[0].view(lowercase_						,	lowercase_						,	height * width).permute(0						,	2						,	1)
          )
          self.assertListEqual(
              list(reshaped_hidden_states.shape[-2:])						,	[num_patches, self.model_tester.embed_dim]						,	)
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  Any    						=		self.model_tester.prepare_config_and_inputs_for_common()

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		(
              self.model_tester.image_size
              if isinstance(self.model_tester.image_size						,	collections.abc.Iterable)
              else (self.model_tester.image_size, self.model_tester.image_size)
          )

          for model_class in self.all_model_classes[:-1]:
              SCREAMING_SNAKE_CASE_      :  List[Any]    						=		True
              self.check_hidden_states_output(lowercase_						,	lowercase_						,	lowercase_						,	lowercase_)

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

              self.check_hidden_states_output(lowercase_						,	lowercase_						,	lowercase_						,	lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self.model_tester.prepare_config_and_inputs_for_common()
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		3

          SCREAMING_SNAKE_CASE_      :  Any    						=		(
              self.model_tester.image_size
              if isinstance(self.model_tester.image_size						,	collections.abc.Iterable)
              else (self.model_tester.image_size, self.model_tester.image_size)
          )
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		(
              config.patch_size
              if isinstance(config.patch_size						,	collections.abc.Iterable)
              else (config.patch_size, config.patch_size)
          )

          SCREAMING_SNAKE_CASE_      :  Any    						=		image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

          for model_class in self.all_model_classes[:-1]:
              SCREAMING_SNAKE_CASE_      :  Any    						=		True
              self.check_hidden_states_output(lowercase_						,	lowercase_						,	lowercase_						,	(padded_height, padded_width))

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              SCREAMING_SNAKE_CASE_      :  Dict    						=		True
              self.check_hidden_states_output(lowercase_						,	lowercase_						,	lowercase_						,	(padded_height, padded_width))
      @slow
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		FocalNetModel.from_pretrained(lowercase_)
              self.assertIsNotNone(lowercase_)



      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


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

          SCREAMING_SNAKE_CASE_      :  Dict    						=		_config_zero_init(lowercase_)
          for model_class in self.all_model_classes:
              SCREAMING_SNAKE_CASE_      :  List[str]    						=		model_class(config=lowercase_)
              for name, param in model.named_parameters():
                  if "embeddings" not in name and 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'						,	)




@require_vision
@require_torch
class 							lowerCAmelCase__				(			unittest.TestCase ):






      '''simple docstring'''

      @cached_property
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          return AutoImageProcessor.from_pretrained('''microsoft/focalnet-tiny''') if is_vision_available() else None



      @slow
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Any    						=		FocalNetForImageClassification.from_pretrained('''microsoft/focalnet-tiny''').to(lowercase_)
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		self.default_image_processor

          SCREAMING_SNAKE_CASE_      :  List[str]    						=		Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''')
          SCREAMING_SNAKE_CASE_      :  str    						=		image_processor(images=lowercase_						,	return_tensors='''pt''').to(lowercase_)

          # forward pass
          with torch.no_grad():
              SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		model(**lowercase_)

          # verify the logits
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		torch.Size((1, 1000))
          self.assertEqual(outputs.logits.shape						,	lowercase_)
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		torch.tensor([0.21_66, -0.43_68, 0.21_91]).to(lowercase_)
          self.assertTrue(torch.allclose(outputs.logits[0, :3]						,	lowercase_						,	atol=1e-4))
          self.assertTrue(outputs.logits.argmax(dim=-1).item()						,	281)




@require_torch
class 							lowerCAmelCase__				(			UpperCAmelCase__				,		unittest.TestCase ):






      '''simple docstring'''

      __UpperCamelCase	 =  (FocalNetBackbone,) if is_torch_available() else ()
      __UpperCamelCase	 =  FocalNetConfig

      __UpperCamelCase	 =  False
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		FocalNetModelTester(self)



import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class 		_lowercase						(					snake_case_						,      snake_case_						,      unittest.TestCase       ):
 lowercase				     =   StableDiffusionDiffEditPipeline
 lowercase				     =   TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'height', 'width', 'image'} | {'image_latents'}
 lowercase				     =   TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {'image'} | {'image_latents'}
 lowercase				     =   frozenset(
     []       )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 lowercase				     =   frozenset([]       )

 def 						SCREAMING_SNAKE_CASE__   ( self	:							List[Any]			)   ->						Dict:

  """simple docstring"""





  torch.manual_seed(0			)
  UpperCamelCase_					:						Dict =					UNetaDConditionModel(
      block_out_channels=(3_2, 6_4)   ,      layers_per_block=2   ,      sample_size=3_2   ,      in_channels=4   ,      out_channels=4   ,      down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   ,      up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')   ,      cross_attention_dim=3_2   ,      attention_head_dim=(2, 4)   ,      use_linear_projection=snake_case   ,      )
  UpperCamelCase_					:						Optional[int] =					DDIMScheduler(
      beta_start=0.00085   ,      beta_end=0.012   ,      beta_schedule='scaled_linear'   ,      clip_sample=snake_case   ,      set_alpha_to_one=snake_case   ,      )
  UpperCamelCase_					:						Tuple =					DDIMInverseScheduler(
      beta_start=0.00085   ,      beta_end=0.012   ,      beta_schedule='scaled_linear'   ,      clip_sample=snake_case   ,      set_alpha_to_zero=snake_case   ,      )
  torch.manual_seed(0			)
  UpperCamelCase_					:						List[Any] =					AutoencoderKL(
      block_out_channels=[3_2, 6_4]   ,      in_channels=3   ,      out_channels=3   ,      down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']   ,      up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']   ,      latent_channels=4   ,      sample_size=1_2_8   ,      )
  torch.manual_seed(0			)
  UpperCamelCase_					:						List[Any] =					CLIPTextConfig(
      bos_token_id=0   ,      eos_token_id=2   ,      hidden_size=3_2   ,      intermediate_size=3_7   ,      layer_norm_eps=1e-05   ,      num_attention_heads=4   ,      num_hidden_layers=5   ,      pad_token_id=1   ,      vocab_size=1_0_0_0   ,      hidden_act='gelu'   ,      projection_dim=5_1_2   ,      )
  UpperCamelCase_					:						List[str] =					CLIPTextModel(snake_case			)
  UpperCamelCase_					:						List[str] =					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'			)

  UpperCamelCase_					:						Any =					{
      'unet': unet,
      'scheduler': scheduler,
      'inverse_scheduler': inverse_scheduler,
      'vae': vae,
      'text_encoder': text_encoder,
      'tokenizer': tokenizer,
      'safety_checker': None,
      'feature_extractor': None,
  }

  return components

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]   ,      snake_case	:							Dict   ,      snake_case	:							Tuple=0			)   ->						int:

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					floats_tensor((1, 1_6, 1_6)   ,      rng=random.Random(snake_case			)			).to(snake_case			)
  UpperCamelCase_					:						Optional[int] =					floats_tensor((1, 2, 4, 1_6, 1_6)   ,      rng=random.Random(snake_case			)			).to(snake_case			)
  if str(snake_case			).startswith('mps'			):
   UpperCamelCase_					:						int =					torch.manual_seed(snake_case			)
  else:
   UpperCamelCase_					:						Any =					torch.Generator(device=snake_case			).manual_seed(snake_case			)
  UpperCamelCase_					:						Union[str, Any] =					{
      'prompt': 'a dog and a newt',
      'mask_image': mask,
      'image_latents': latents,
      'generator': generator,
      'num_inference_steps': 2,
      'inpaint_strength': 1.0,
      'guidance_scale': 6.0,
      'output_type': 'numpy',
  }

  return inputs

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Union[str, Any]   ,      snake_case	:							Union[str, Any]   ,      snake_case	:							Dict=0			)   ->						str:

  """simple docstring"""





  UpperCamelCase_					:						List[Any] =					floats_tensor((1, 3, 3_2, 3_2)   ,      rng=random.Random(snake_case			)			).to(snake_case			)
  UpperCamelCase_					:						int =					image.cpu().permute(0   ,      2   ,      3   ,      1			)[0]
  UpperCamelCase_					:						Optional[int] =					Image.fromarray(np.uinta(snake_case			)			).convert('RGB'			)
  if str(snake_case			).startswith('mps'			):
   UpperCamelCase_					:						int =					torch.manual_seed(snake_case			)
  else:
   UpperCamelCase_					:						Optional[int] =					torch.Generator(device=snake_case			).manual_seed(snake_case			)
  UpperCamelCase_					:						List[str] =					{
      'image': image,
      'source_prompt': 'a cat and a frog',
      'target_prompt': 'a dog and a newt',
      'generator': generator,
      'num_inference_steps': 2,
      'num_maps_per_mask': 2,
      'mask_encode_strength': 1.0,
      'guidance_scale': 6.0,
      'output_type': 'numpy',
  }

  return inputs

 def 						SCREAMING_SNAKE_CASE__   ( self	:							List[str]   ,      snake_case	:							str   ,      snake_case	:							Dict=0			)   ->						int:

  """simple docstring"""





  UpperCamelCase_					:						List[Any] =					floats_tensor((1, 3, 3_2, 3_2)   ,      rng=random.Random(snake_case			)			).to(snake_case			)
  UpperCamelCase_					:						Any =					image.cpu().permute(0   ,      2   ,      3   ,      1			)[0]
  UpperCamelCase_					:						List[str] =					Image.fromarray(np.uinta(snake_case			)			).convert('RGB'			)
  if str(snake_case			).startswith('mps'			):
   UpperCamelCase_					:						Any =					torch.manual_seed(snake_case			)
  else:
   UpperCamelCase_					:						int =					torch.Generator(device=snake_case			).manual_seed(snake_case			)
  UpperCamelCase_					:						int =					{
      'image': image,
      'prompt': 'a cat and a frog',
      'generator': generator,
      'num_inference_steps': 2,
      'inpaint_strength': 1.0,
      'guidance_scale': 6.0,
      'decode_latents': True,
      'output_type': 'numpy',
  }
  return inputs

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

  """simple docstring"""





  if not hasattr(self.pipeline_class   ,      '_optional_components'			):
   return

  UpperCamelCase_					:						int =					self.get_dummy_components()
  UpperCamelCase_					:						str =					self.pipeline_class(**snake_case			)
  pipe.to(snake_case			)
  pipe.set_progress_bar_config(disable=snake_case			)

  # set all optional components to None and update pipeline config accordingly
  for optional_component in pipe._optional_components:
   setattr(snake_case   ,      snake_case   ,      snake_case			)
  pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}			)

  UpperCamelCase_					:						Any =					self.get_dummy_inputs(snake_case			)
  UpperCamelCase_					:						Any =					pipe(**snake_case			)[0]

  with tempfile.TemporaryDirectory() as tmpdir:
   pipe.save_pretrained(snake_case			)
   UpperCamelCase_					:						int =					self.pipeline_class.from_pretrained(snake_case			)
   pipe_loaded.to(snake_case			)
   pipe_loaded.set_progress_bar_config(disable=snake_case			)

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

  UpperCamelCase_					:						Dict =					self.get_dummy_inputs(snake_case			)
  UpperCamelCase_					:						int =					pipe_loaded(**snake_case			)[0]

  UpperCamelCase_					:						int =					np.abs(output - output_loaded			).max()
  self.assertLess(snake_case   ,      1e-4			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							str			)   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						Optional[int] =					'cpu'

  UpperCamelCase_					:						Any =					self.get_dummy_components()
  UpperCamelCase_					:						int =					self.pipeline_class(**snake_case			)
  pipe.to(snake_case			)
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						Union[str, Any] =					self.get_dummy_mask_inputs(snake_case			)
  UpperCamelCase_					:						Optional[int] =					pipe.generate_mask(**snake_case			)
  UpperCamelCase_					:						List[str] =					mask[0, -3:, -3:]

  self.assertEqual(mask.shape   ,      (1, 1_6, 1_6)			)
  UpperCamelCase_					:						Tuple =					np.array([0] * 9			)
  UpperCamelCase_					:						int =					np.abs(mask_slice.flatten() - expected_slice			).max()
  self.assertLessEqual(snake_case   ,      1e-3			)
  self.assertEqual(mask[0, -3, -4]   ,      0			)

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

  """simple docstring"""





  UpperCamelCase_					:						Dict =					'cpu'

  UpperCamelCase_					:						List[str] =					self.get_dummy_components()
  UpperCamelCase_					:						List[Any] =					self.pipeline_class(**snake_case			)
  pipe.to(snake_case			)
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						Optional[int] =					self.get_dummy_inversion_inputs(snake_case			)
  UpperCamelCase_					:						Optional[Any] =					pipe.invert(**snake_case			).images
  UpperCamelCase_					:						Tuple =					image[0, -1, -3:, -3:]

  self.assertEqual(image.shape   ,      (2, 3_2, 3_2, 3)			)
  UpperCamelCase_					:						str =					np.array(
      [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.51050, 0.5015, 0.4407, 0.4799]   ,      )
  UpperCamelCase_					:						Any =					np.abs(image_slice.flatten() - expected_slice			).max()
  self.assertLessEqual(snake_case   ,      1e-3			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any			)   ->						Any:

  """simple docstring"""





  super().test_inference_batch_single_identical(expected_max_diff=5e-3			)


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

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					'cpu'

  UpperCamelCase_					:						Tuple =					self.get_dummy_components()

  UpperCamelCase_					:						int =					{'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': 'scaled_linear'}
  UpperCamelCase_					:						str =					DPMSolverMultistepScheduler(**snake_case			)
  UpperCamelCase_					:						Union[str, Any] =					DPMSolverMultistepInverseScheduler(**snake_case			)

  UpperCamelCase_					:						Optional[int] =					self.pipeline_class(**snake_case			)
  pipe.to(snake_case			)
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						Any =					self.get_dummy_inversion_inputs(snake_case			)
  UpperCamelCase_					:						Union[str, Any] =					pipe.invert(**snake_case			).images
  UpperCamelCase_					:						Any =					image[0, -1, -3:, -3:]

  self.assertEqual(image.shape   ,      (2, 3_2, 3_2, 3)			)
  UpperCamelCase_					:						int =					np.array(
      [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.51050, 0.5015, 0.4407, 0.4799]   ,      )
  UpperCamelCase_					:						Union[str, Any] =					np.abs(image_slice.flatten() - expected_slice			).max()
  self.assertLessEqual(snake_case   ,      1e-3			)





@require_torch_gpu
@slow
class 		_lowercase						(					unittest.TestCase       ):

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						Any:

  """simple docstring"""





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

 @classmethod
 def 						SCREAMING_SNAKE_CASE__   ( cls	:							Any			)   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						Optional[int] =					load_image(
      'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'			)

  UpperCamelCase_					:						Any =					raw_image.convert('RGB'			).resize((7_6_8, 7_6_8)			)

  UpperCamelCase_					:						Optional[Any] =					raw_image

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

  """simple docstring"""





  UpperCamelCase_					:						Dict =					torch.manual_seed(0			)

  UpperCamelCase_					:						List[str] =					StableDiffusionDiffEditPipeline.from_pretrained(
      'stabilityai/stable-diffusion-2-1'   ,      safety_checker=snake_case   ,      torch_dtype=torch.floataa			)
  UpperCamelCase_					:						int =					DDIMScheduler.from_config(pipe.scheduler.config			)
  UpperCamelCase_					:						Tuple =					DDIMInverseScheduler.from_config(pipe.scheduler.config			)
  pipe.enable_model_cpu_offload()
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						Optional[int] =					'a bowl of fruit'
  UpperCamelCase_					:						int =					'a bowl of pears'

  UpperCamelCase_					:						Any =					pipe.generate_mask(
      image=self.raw_image   ,      source_prompt=snake_case   ,      target_prompt=snake_case   ,      generator=snake_case   ,      )

  UpperCamelCase_					:						int =					pipe.invert(
      prompt=snake_case   ,      image=self.raw_image   ,      inpaint_strength=0.7   ,      generator=snake_case			).latents

  UpperCamelCase_					:						Optional[int] =					pipe(
      prompt=snake_case   ,      mask_image=snake_case   ,      image_latents=snake_case   ,      generator=snake_case   ,      negative_prompt=snake_case   ,      inpaint_strength=0.7   ,      output_type='numpy'   ,      ).images[0]

  UpperCamelCase_					:						Any =					(
      np.array(
          load_image(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/diffedit/pears.png'			).resize((7_6_8, 7_6_8)			)			)
      / 2_5_5
  )
  assert np.abs((expected_image - image).max()			) < 5e-1


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

  """simple docstring"""





  UpperCamelCase_					:						Union[str, Any] =					torch.manual_seed(0			)

  UpperCamelCase_					:						List[Any] =					StableDiffusionDiffEditPipeline.from_pretrained(
      'stabilityai/stable-diffusion-2-1'   ,      safety_checker=snake_case   ,      torch_dtype=torch.floataa			)
  UpperCamelCase_					:						Optional[int] =					DPMSolverMultistepScheduler.from_config(pipe.scheduler.config			)
  UpperCamelCase_					:						Union[str, Any] =					DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config			)
  pipe.enable_model_cpu_offload()
  pipe.set_progress_bar_config(disable=snake_case			)

  UpperCamelCase_					:						str =					'a bowl of fruit'
  UpperCamelCase_					:						int =					'a bowl of pears'

  UpperCamelCase_					:						Union[str, Any] =					pipe.generate_mask(
      image=self.raw_image   ,      source_prompt=snake_case   ,      target_prompt=snake_case   ,      generator=snake_case   ,      )

  UpperCamelCase_					:						Any =					pipe.invert(
      prompt=snake_case   ,      image=self.raw_image   ,      inpaint_strength=0.7   ,      generator=snake_case   ,      num_inference_steps=2_5   ,      ).latents

  UpperCamelCase_					:						int =					pipe(
      prompt=snake_case   ,      mask_image=snake_case   ,      image_latents=snake_case   ,      generator=snake_case   ,      negative_prompt=snake_case   ,      inpaint_strength=0.7   ,      num_inference_steps=2_5   ,      output_type='numpy'   ,      ).images[0]

  UpperCamelCase_					:						str =					(
      np.array(
          load_image(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/diffedit/pears.png'			).resize((7_6_8, 7_6_8)			)			)
      / 2_5_5
  )
  assert np.abs((expected_image - image).max()			) < 5e-1



