



"""simple docstring"""


from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
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 import (
				    TFDeiTForImageClassification,
				    TFDeiTForImageClassificationWithTeacher,
				    TFDeiTForMaskedImageModeling,
				    TFDeiTModel,
				)
				from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
				from PIL import Image

				from transformers import DeiTImageProcessor







class 				A_      :



						"""simple docstring"""






						def __init__( self			:Dict       ,  lowercase_			:Dict       ,  lowercase_			:Dict=13       ,  lowercase_			:Any=30       ,  lowercase_			:List[Any]=2       ,  lowercase_			:Optional[Any]=3       ,  lowercase_			:Any=True       ,  lowercase_			:str=True       ,  lowercase_			:int=32       ,  lowercase_			:Optional[int]=2       ,  lowercase_			:Union[str, Any]=4       ,  lowercase_			:Dict=37       ,  lowercase_			:Any="gelu"       ,  lowercase_			:Tuple=0.1       ,  lowercase_			:Union[str, Any]=0.1       ,  lowercase_			:List[str]=10       ,  lowercase_			:Dict=0.02       ,  lowercase_			:str=3       ,  lowercase_			:Dict=None       ,  lowercase_			:Any=2       ,  )    ->      Optional[Any]:
												UpperCAmelCase					=			parent
												UpperCAmelCase					=			batch_size
												UpperCAmelCase					=			image_size
												UpperCAmelCase					=			patch_size
												UpperCAmelCase					=			num_channels
												UpperCAmelCase					=			is_training
												UpperCAmelCase					=			use_labels
												UpperCAmelCase					=			hidden_size
												UpperCAmelCase					=			num_hidden_layers
												UpperCAmelCase					=			num_attention_heads
												UpperCAmelCase					=			intermediate_size
												UpperCAmelCase					=			hidden_act
												UpperCAmelCase					=			hidden_dropout_prob
												UpperCAmelCase					=			attention_probs_dropout_prob
												UpperCAmelCase					=			type_sequence_label_size
												UpperCAmelCase					=			initializer_range
												UpperCAmelCase					=			scope
												UpperCAmelCase					=			encoder_stride

												# in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
												UpperCAmelCase					=			(image_size // patch_size) ** 2
												UpperCAmelCase					=			num_patches + 2






						def      UpperCAmelCase__					( self			:Any )    ->      Optional[Any]:
												UpperCAmelCase					=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

												UpperCAmelCase					=			None
												if self.use_labels:
																		UpperCAmelCase					=			ids_tensor([self.batch_size]       ,  self.type_sequence_label_size )

												UpperCAmelCase					=			self.get_config()

												return config, pixel_values, labels






						def      UpperCAmelCase__					( self			:Optional[int] )    ->      Tuple:
												return DeiTConfig(
												    image_size=self.image_size       ,  patch_size=self.patch_size       ,  num_channels=self.num_channels       ,  hidden_size=self.hidden_size       ,  num_hidden_layers=self.num_hidden_layers       ,  num_attention_heads=self.num_attention_heads       ,  intermediate_size=self.intermediate_size       ,  hidden_act=self.hidden_act       ,  hidden_dropout_prob=self.hidden_dropout_prob       ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,  is_decoder=_SCREAMING_SNAKE_CASE       ,  initializer_range=self.initializer_range       ,  encoder_stride=self.encoder_stride       ,  )






						def      UpperCAmelCase__					( self			:List[Any]       ,  lowercase_			:Tuple       ,  lowercase_			:Dict       ,  lowercase_			:str )    ->      Tuple:
												UpperCAmelCase					=			TFDeiTModel(config=_SCREAMING_SNAKE_CASE )
												UpperCAmelCase					=			model(_SCREAMING_SNAKE_CASE )
												self.parent.assertEqual(result.last_hidden_state.shape       ,  (self.batch_size, self.seq_length, self.hidden_size) )






						def      UpperCAmelCase__					( self			:int       ,  lowercase_			:Union[str, Any]       ,  lowercase_			:int       ,  lowercase_			:Optional[Any] )    ->      Dict:
												UpperCAmelCase					=			TFDeiTForMaskedImageModeling(config=_SCREAMING_SNAKE_CASE )
												UpperCAmelCase					=			model(_SCREAMING_SNAKE_CASE )
												self.parent.assertEqual(
												    result.reconstruction.shape       ,  (self.batch_size, self.num_channels, self.image_size, self.image_size) )

												# test greyscale images
												UpperCAmelCase					=			1
												UpperCAmelCase					=			TFDeiTForMaskedImageModeling(_SCREAMING_SNAKE_CASE )

												UpperCAmelCase					=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
												UpperCAmelCase					=			model(_SCREAMING_SNAKE_CASE )
												self.parent.assertEqual(result.reconstruction.shape       ,  (self.batch_size, 1, self.image_size, self.image_size) )






						def      UpperCAmelCase__					( self			:Optional[int]       ,  lowercase_			:Dict       ,  lowercase_			:Dict       ,  lowercase_			:List[Any] )    ->      Optional[int]:
												UpperCAmelCase					=			self.type_sequence_label_size
												UpperCAmelCase					=			TFDeiTForImageClassification(_SCREAMING_SNAKE_CASE )
												UpperCAmelCase					=			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
												UpperCAmelCase					=			1
												UpperCAmelCase					=			TFDeiTForImageClassification(_SCREAMING_SNAKE_CASE )

												UpperCAmelCase					=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
												UpperCAmelCase					=			model(_SCREAMING_SNAKE_CASE       ,  labels=_SCREAMING_SNAKE_CASE )
												self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.type_sequence_label_size) )






						def      UpperCAmelCase__					( self			:Tuple )    ->      int:
												UpperCAmelCase					=			self.prepare_config_and_inputs()
												UpperCAmelCase					=			config_and_inputs
												UpperCAmelCase					=			{"pixel_values": pixel_values}
												return config, inputs_dict







@require_tf
class 				A_      (					__SCREAMING_SNAKE_CASE     ,       __SCREAMING_SNAKE_CASE     ,       unittest.TestCase    ):



						"""simple docstring"""
						__UpperCamelCase		    =							(
						    (
						        TFDeiTModel,
						        TFDeiTForImageClassification,
						        TFDeiTForImageClassificationWithTeacher,
						        TFDeiTForMaskedImageModeling,
						    )
						    if is_tf_available()
						    else ()
						)
						__UpperCamelCase		    =							(
						    {
						        """feature-extraction""": TFDeiTModel,
						        """image-classification""": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
						    }
						    if is_tf_available()
						    else {}
						)

						__UpperCamelCase		    =							False
						__UpperCamelCase		    =							False
						__UpperCamelCase		    =							False
						__UpperCamelCase		    =							False






						def      UpperCAmelCase__					( self			:List[Any] )    ->      Dict:
												UpperCAmelCase					=			TFDeiTModelTester(self )
												UpperCAmelCase					=			ConfigTester(self       ,  config_class=_SCREAMING_SNAKE_CASE       ,  has_text_modality=_SCREAMING_SNAKE_CASE       ,  hidden_size=37 )






						def      UpperCAmelCase__					( self			:Tuple )    ->      int:
												self.config_tester.run_common_tests()






						@unittest.skip(reason='DeiT does not use inputs_embeds' )
						def      UpperCAmelCase__					( self			:Union[str, Any] )    ->      Dict:
												pass






						def      UpperCAmelCase__					( self			:str )    ->      Dict:
												UpperCAmelCase					=			self.model_tester.prepare_config_and_inputs_for_common()

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






						def      UpperCAmelCase__					( self			:Any )    ->      Any:
												UpperCAmelCase					=			self.model_tester.prepare_config_and_inputs_for_common()

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

																		UpperCAmelCase					=			["pixel_values"]
																		self.assertListEqual(arg_names[:1]       ,  _SCREAMING_SNAKE_CASE )






						def      UpperCAmelCase__					( self			:str )    ->      int:
												UpperCAmelCase					=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE )






						def      UpperCAmelCase__					( self			:Any )    ->      Dict:
												UpperCAmelCase					=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_image_modeling(*_SCREAMING_SNAKE_CASE )






						def      UpperCAmelCase__					( self			:Optional[Any] )    ->      Union[str, Any]:
												UpperCAmelCase					=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE )






						def      UpperCAmelCase__					( self			:List[Any]       ,  lowercase_			:Dict       ,  lowercase_			:Tuple       ,  lowercase_			:Union[str, Any]=False )    ->      str:
												UpperCAmelCase					=			super()._prepare_for_class(_SCREAMING_SNAKE_CASE       ,  _SCREAMING_SNAKE_CASE       ,  return_labels=_SCREAMING_SNAKE_CASE )

												if return_labels:
																		if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call ).parameters:
																								del inputs_dict["labels"]

												return inputs_dict






						@slow
						def      UpperCAmelCase__					( self			:Optional[Any] )    ->      Dict:

												for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		UpperCAmelCase					=			TFDeiTModel.from_pretrained(_SCREAMING_SNAKE_CASE )
																		self.assertIsNotNone(_SCREAMING_SNAKE_CASE )


def 			_lowerCAmelCase				(       ):
						UpperCAmelCase					=			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
						return image







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



						"""simple docstring"""






						@cached_property
						def      UpperCAmelCase__					( self			:Tuple )    ->      Any:
												return (
												    DeiTImageProcessor.from_pretrained('facebook/deit-base-distilled-patch16-224' )
												    if is_vision_available()
												    else None
												)






						@slow
						def      UpperCAmelCase__					( self			:Any )    ->      List[Any]:
												UpperCAmelCase					=			TFDeiTForImageClassificationWithTeacher.from_pretrained('facebook/deit-base-distilled-patch16-224' )

												UpperCAmelCase					=			self.default_image_processor
												UpperCAmelCase					=			prepare_img()
												UpperCAmelCase					=			image_processor(images=_SCREAMING_SNAKE_CASE       ,  return_tensors='tf' )

												# forward pass
												UpperCAmelCase					=			model(**_SCREAMING_SNAKE_CASE )

												# verify the logits
												UpperCAmelCase					=			tf.TensorShape((1, 10_00) )
												self.assertEqual(outputs.logits.shape       ,  _SCREAMING_SNAKE_CASE )

												UpperCAmelCase					=			tf.constant([-1.0266, 0.1912, -1.2861] )

												self.assertTrue(np.allclose(outputs.logits[0, :3]       ,  _SCREAMING_SNAKE_CASE       ,  atol=1E-4 ) )










"""simple docstring"""


from typing import Callable, List, Optional, Union

import PIL
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPSegForImageSegmentation,
    CLIPSegProcessor,
    CLIPTextModel,
    CLIPTokenizer,
)

from diffusers import DiffusionPipeline
from diffusers.configuration_utils import FrozenDict
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import deprecate, is_accelerate_available, logging


snake_case_    					=	logging.get_logger(__name__)  # pylint: disable=invalid-name







class 				A_      (					SCREAMING_SNAKE_CASE_    ):



						"""simple docstring"""






						def __init__( self			:Optional[Any]       ,  lowercase_			:CLIPSegForImageSegmentation       ,  lowercase_			:CLIPSegProcessor       ,  lowercase_			:AutoencoderKL       ,  lowercase_			:CLIPTextModel       ,  lowercase_			:CLIPTokenizer       ,  lowercase_			:UNetaDConditionModel       ,  lowercase_			:Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]       ,  lowercase_			:StableDiffusionSafetyChecker       ,  lowercase_			:CLIPImageProcessor       ,  )    ->      List[str]:
												super().__init__()

												if hasattr(scheduler.config       ,  'steps_offset' ) and scheduler.config.steps_offset != 1:
																		UpperCAmelCase					=			(
																		    f"""The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"""
																		    f""" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure """
																		    'to update the config accordingly as leaving `steps_offset` might led to incorrect results'
																		    ' in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,'
																		    ' it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`'
																		    ' file'
																		)
																		deprecate('steps_offset!=1'       ,  '1.0.0'       ,  lowercase_       ,  standard_warn=lowercase_ )
																		UpperCAmelCase					=			dict(scheduler.config )
																		UpperCAmelCase					=			1
																		UpperCAmelCase					=			FrozenDict(lowercase_ )

												if hasattr(scheduler.config       ,  'skip_prk_steps' ) and scheduler.config.skip_prk_steps is False:
																		UpperCAmelCase					=			(
																		    f"""The configuration file of this scheduler: {scheduler} has not set the configuration"""
																		    ' `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make'
																		    ' sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to'
																		    ' incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face'
																		    ' Hub, it would be very nice if you could open a Pull request for the'
																		    ' `scheduler/scheduler_config.json` file'
																		)
																		deprecate('skip_prk_steps not set'       ,  '1.0.0'       ,  lowercase_       ,  standard_warn=lowercase_ )
																		UpperCAmelCase					=			dict(scheduler.config )
																		UpperCAmelCase					=			True
																		UpperCAmelCase					=			FrozenDict(lowercase_ )

												if safety_checker is None:
																		logger.warning(
																		    f"""You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"""
																		    ' that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered'
																		    ' results in services or applications open to the public. Both the diffusers team and Hugging Face'
																		    ' strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling'
																		    ' it only for use-cases that involve analyzing network behavior or auditing its results. For more'
																		    ' information, please have a look at https://github.com/huggingface/diffusers/pull/254 .' )

												self.register_modules(
												    segmentation_model=lowercase_       ,  segmentation_processor=lowercase_       ,  vae=lowercase_       ,  text_encoder=lowercase_       ,  tokenizer=lowercase_       ,  unet=lowercase_       ,  scheduler=lowercase_       ,  safety_checker=lowercase_       ,  feature_extractor=lowercase_       ,  )






						def      UpperCAmelCase__					( self			:Optional[int]       ,  lowercase_			:Optional[Union[str, int]] = "auto" )    ->      int:
												if slice_size == "auto":
																		# half the attention head size is usually a good trade-off between
																		# speed and memory
																		UpperCAmelCase					=			self.unet.config.attention_head_dim // 2
												self.unet.set_attention_slice(lowercase_ )






						def      UpperCAmelCase__					( self			:Optional[int] )    ->      List[str]:
												self.enable_attention_slicing(lowercase_ )






						def      UpperCAmelCase__					( self			:int )    ->      int:
												if is_accelerate_available():
																		from accelerate import cpu_offload
												else:
																		raise ImportError('Please install accelerate via `pip install accelerate`' )

												UpperCAmelCase					=			torch.device('cuda' )

												for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae, self.safety_checker]:
																		if cpu_offloaded_model is not None:
																								cpu_offload(lowercase_       ,  lowercase_ )






						@property
						# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
						def      UpperCAmelCase__					( self			:Optional[Any] )    ->      List[str]:
												if self.device != torch.device('meta' ) or not hasattr(self.unet       ,  '_hf_hook' ):
																		return self.device
												for module in self.unet.modules():
																		if (
																		    hasattr(lowercase_       ,  '_hf_hook' )
																		    and hasattr(module._hf_hook       ,  'execution_device' )
																		    and module._hf_hook.execution_device is not None
																		):
																								return torch.device(module._hf_hook.execution_device )
												return self.device






						@torch.no_grad()
						def __call__( self			:Optional[Any]       ,  lowercase_			:Union[str, List[str]]       ,  lowercase_			:Union[torch.FloatTensor, PIL.Image.Image]       ,  lowercase_			:str       ,  lowercase_			:int = 5_12       ,  lowercase_			:int = 5_12       ,  lowercase_			:int = 50       ,  lowercase_			:float = 7.5       ,  lowercase_			:Optional[Union[str, List[str]]] = None       ,  lowercase_			:Optional[int] = 1       ,  lowercase_			:float = 0.0       ,  lowercase_			:Optional[torch.Generator] = None       ,  lowercase_			:Optional[torch.FloatTensor] = None       ,  lowercase_			:Optional[str] = "pil"       ,  lowercase_			:bool = True       ,  lowercase_			:Optional[Callable[[int, int, torch.FloatTensor], None]] = None       ,  lowercase_			:int = 1       ,  **lowercase_			:int       ,  )    ->      int:
												UpperCAmelCase					=			self.segmentation_processor(
												    text=[text]       ,  images=[image]       ,  padding='max_length'       ,  return_tensors='pt' ).to(self.device )
												UpperCAmelCase					=			self.segmentation_model(**lowercase_ )
												UpperCAmelCase					=			torch.sigmoid(outputs.logits ).cpu().detach().unsqueeze(-1 ).numpy()
												UpperCAmelCase					=			self.numpy_to_pil(lowercase_ )[0].resize(image.size )

												# Run inpainting pipeline with the generated mask
												UpperCAmelCase					=			StableDiffusionInpaintPipeline(
												    vae=self.vae       ,  text_encoder=self.text_encoder       ,  tokenizer=self.tokenizer       ,  unet=self.unet       ,  scheduler=self.scheduler       ,  safety_checker=self.safety_checker       ,  feature_extractor=self.feature_extractor       ,  )
												return inpainting_pipeline(
												    prompt=lowercase_       ,  image=lowercase_       ,  mask_image=lowercase_       ,  height=lowercase_       ,  width=lowercase_       ,  num_inference_steps=lowercase_       ,  guidance_scale=lowercase_       ,  negative_prompt=lowercase_       ,  num_images_per_prompt=lowercase_       ,  eta=lowercase_       ,  generator=lowercase_       ,  latents=lowercase_       ,  output_type=lowercase_       ,  return_dict=lowercase_       ,  callback=lowercase_       ,  callback_steps=lowercase_       ,  )





