



'''simple docstring'''

import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
	import torch
	from torch import nn

	from transformers import (
	    MODEL_MAPPING,
	    BeitForImageClassification,
	    BeitForMaskedImageModeling,
	    BeitForSemanticSegmentation,
	    BeitModel,
	)
	from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	import PIL
	from PIL import Image

	from transformers import BeitImageProcessor


class 	a__     :


	def __init__(self : List[Any],					__UpperCAmelCase : List[str],					__UpperCAmelCase : Optional[Any]=100,					__UpperCAmelCase : str=13,					__UpperCAmelCase : Optional[Any]=30,					__UpperCAmelCase : List[Any]=2,					__UpperCAmelCase : Dict=3,					__UpperCAmelCase : Optional[int]=True,					__UpperCAmelCase : Optional[int]=True,					__UpperCAmelCase : Optional[int]=32,					__UpperCAmelCase : Tuple=4,					__UpperCAmelCase : Any=4,					__UpperCAmelCase : List[str]=37,					__UpperCAmelCase : Optional[int]="gelu",					__UpperCAmelCase : Optional[int]=0.1,					__UpperCAmelCase : List[Any]=0.1,					__UpperCAmelCase : List[Any]=10,					__UpperCAmelCase : List[str]=0.02,					__UpperCAmelCase : Tuple=3,					__UpperCAmelCase : Optional[Any]=None,					__UpperCAmelCase : List[Any]=[0, 1, 2, 3],					)     ->			Union[str, Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= parent
					SCREAMING_SNAKE_CASE  :						Tuple      					= 100
					SCREAMING_SNAKE_CASE  :						str      					= batch_size
					SCREAMING_SNAKE_CASE  :						Dict      					= image_size
					SCREAMING_SNAKE_CASE  :						List[str]      					= patch_size
					SCREAMING_SNAKE_CASE  :						List[Any]      					= num_channels
					SCREAMING_SNAKE_CASE  :						List[str]      					= is_training
					SCREAMING_SNAKE_CASE  :						List[str]      					= use_labels
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= hidden_size
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= num_hidden_layers
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= num_attention_heads
					SCREAMING_SNAKE_CASE  :						str      					= intermediate_size
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= hidden_act
					SCREAMING_SNAKE_CASE  :						Any      					= hidden_dropout_prob
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= attention_probs_dropout_prob
					SCREAMING_SNAKE_CASE  :						Any      					= type_sequence_label_size
					SCREAMING_SNAKE_CASE  :						Tuple      					= initializer_range
					SCREAMING_SNAKE_CASE  :						List[Any]      					= scope
					SCREAMING_SNAKE_CASE  :						Tuple      					= out_indices
					SCREAMING_SNAKE_CASE  :						Tuple      					= num_labels

					# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
					SCREAMING_SNAKE_CASE  :						str      					= (image_size // patch_size) ** 2
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= num_patches + 1


	def 				lowercase__  (self : Any      )     ->			str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

					SCREAMING_SNAKE_CASE  :						Dict      					= None
					SCREAMING_SNAKE_CASE  :						List[Any]      					= None
					if self.use_labels:
									SCREAMING_SNAKE_CASE  :						List[Any]      					= ids_tensor([self.batch_size],					self.type_sequence_label_size      )
									SCREAMING_SNAKE_CASE  :						Dict      					= ids_tensor([self.batch_size, self.image_size, self.image_size],					self.num_labels      )

					SCREAMING_SNAKE_CASE  :						List[Any]      					= self.get_config()

					return config, pixel_values, labels, pixel_labels


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

					"""simple docstring"""

					return BeitConfig(
					    vocab_size=self.vocab_size,					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=__UpperCAmelCase,					initializer_range=self.initializer_range,					out_indices=self.out_indices,					)


	def 				lowercase__  (self : List[str],					__UpperCAmelCase : Union[str, Any],					__UpperCAmelCase : List[Any],					__UpperCAmelCase : List[str],					__UpperCAmelCase : Optional[int]      )     ->			Union[str, Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						int      					= BeitModel(config=__UpperCAmelCase      )
					model.to(__UpperCAmelCase      )
					model.eval()
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= model(__UpperCAmelCase      )
					self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.seq_length, self.hidden_size)      )


	def 				lowercase__  (self : List[Any],					__UpperCAmelCase : Dict,					__UpperCAmelCase : List[Any],					__UpperCAmelCase : Optional[Any],					__UpperCAmelCase : Union[str, Any]      )     ->			int:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= BeitForMaskedImageModeling(config=__UpperCAmelCase      )
					model.to(__UpperCAmelCase      )
					model.eval()
					SCREAMING_SNAKE_CASE  :						Any      					= model(__UpperCAmelCase      )
					self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.seq_length - 1, self.vocab_size)      )


	def 				lowercase__  (self : Optional[Any],					__UpperCAmelCase : Optional[Any],					__UpperCAmelCase : Dict,					__UpperCAmelCase : Union[str, Any],					__UpperCAmelCase : List[Any]      )     ->			str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= self.type_sequence_label_size
					SCREAMING_SNAKE_CASE  :						int      					= BeitForImageClassification(__UpperCAmelCase      )
					model.to(__UpperCAmelCase      )
					model.eval()
					SCREAMING_SNAKE_CASE  :						List[Any]      					= model(__UpperCAmelCase,					labels=__UpperCAmelCase      )
					self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.type_sequence_label_size)      )

					# test greyscale images
					SCREAMING_SNAKE_CASE  :						str      					= 1
					SCREAMING_SNAKE_CASE  :						Any      					= BeitForImageClassification(__UpperCAmelCase      )
					model.to(__UpperCAmelCase      )
					model.eval()

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
					SCREAMING_SNAKE_CASE  :						List[str]      					= model(__UpperCAmelCase,					labels=__UpperCAmelCase      )
					self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.type_sequence_label_size)      )


	def 				lowercase__  (self : List[Any],					__UpperCAmelCase : str,					__UpperCAmelCase : Optional[Any],					__UpperCAmelCase : int,					__UpperCAmelCase : List[Any]      )     ->			Union[str, Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= self.num_labels
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= BeitForSemanticSegmentation(__UpperCAmelCase      )
					model.to(__UpperCAmelCase      )
					model.eval()
					SCREAMING_SNAKE_CASE  :						Tuple      					= model(__UpperCAmelCase      )
					self.parent.assertEqual(
					    result.logits.shape,					(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)      )
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= model(__UpperCAmelCase,					labels=__UpperCAmelCase      )
					self.parent.assertEqual(
					    result.logits.shape,					(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)      )






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

					"""simple docstring"""

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



@require_torch
class 	a__     (				_lowercase,			_lowercase,			unittest.TestCase	):
	__magic_name__		:    Union[str, Any]	    =     (
	    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
	    if is_torch_available()
	    else ()
	)
	__magic_name__		:    Tuple	    =     (
	    {
	        "feature-extraction": BeitModel,
	        "image-classification": BeitForImageClassification,
	        "image-segmentation": BeitForSemanticSegmentation,
	    }
	    if is_torch_available()
	    else {}
	)

	__magic_name__		:    int	    =     False
	__magic_name__		:    Tuple	    =     False
	__magic_name__		:    int	    =     False


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= BeitModelTester(self      )
					SCREAMING_SNAKE_CASE  :						Any      					= ConfigTester(self,					config_class=__UpperCAmelCase,					has_text_modality=__UpperCAmelCase,					hidden_size=37      )


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

					"""simple docstring"""

					self.config_tester.run_common_tests()


	@unittest.skip(reason='''BEiT does not use inputs_embeds'''      )
	def 				lowercase__  (self : Tuple      )     ->			Tuple:

					"""simple docstring"""

					pass


	@require_torch_multi_gpu
	@unittest.skip(reason='''BEiT has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''      )
	def 				lowercase__  (self : str      )     ->			Any:

					"""simple docstring"""

					pass


	def 				lowercase__  (self : Any      )     ->			Dict:

					"""simple docstring"""

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

					for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE  :						Dict      					= model_class(__UpperCAmelCase      )
									self.assertIsInstance(model.get_input_embeddings(),					(nn.Module)      )
									SCREAMING_SNAKE_CASE  :						str      					= model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(__UpperCAmelCase,					nn.Linear      )      )


	def 				lowercase__  (self : str      )     ->			Dict:

					"""simple docstring"""

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

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

									SCREAMING_SNAKE_CASE  :						Any      					= ['''pixel_values''']
									self.assertListEqual(arg_names[:1],					__UpperCAmelCase      )


	def 				lowercase__  (self : Optional[int]      )     ->			Dict:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Tuple      					= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*__UpperCAmelCase      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase      )


	def 				lowercase__  (self : Tuple      )     ->			Tuple:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_semantic_segmentation(*__UpperCAmelCase      )


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

					"""simple docstring"""

					if not self.model_tester.is_training:
									return

					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE  :						Optional[int]      					= self.model_tester.prepare_config_and_inputs_for_common()
					SCREAMING_SNAKE_CASE  :						Dict      					= True

					for model_class in self.all_model_classes:
									# we don't test BeitForMaskedImageModeling
									if model_class in [*get_values(__UpperCAmelCase      ), BeitForMaskedImageModeling]:
													continue

									SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= model_class(__UpperCAmelCase      )
									model.to(__UpperCAmelCase      )
									model.train()
									SCREAMING_SNAKE_CASE  :						List[str]      					= self._prepare_for_class(__UpperCAmelCase,					__UpperCAmelCase,					return_labels=__UpperCAmelCase      )
									SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= model(**__UpperCAmelCase      ).loss
									loss.backward()


	def 				lowercase__  (self : Any      )     ->			List[str]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE  :						Optional[int]      					= self.model_tester.prepare_config_and_inputs_for_common()
					if not self.model_tester.is_training:
									return

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= False
					SCREAMING_SNAKE_CASE  :						int      					= True

					for model_class in self.all_model_classes:
									# we don't test BeitForMaskedImageModeling
									if (
									    model_class in [*get_values(__UpperCAmelCase      ), BeitForMaskedImageModeling]
									    or not model_class.supports_gradient_checkpointing
									):
													continue

									SCREAMING_SNAKE_CASE  :						List[str]      					= model_class(__UpperCAmelCase      )
									model.gradient_checkpointing_enable()
									model.to(__UpperCAmelCase      )
									model.train()
									SCREAMING_SNAKE_CASE  :						List[str]      					= self._prepare_for_class(__UpperCAmelCase,					__UpperCAmelCase,					return_labels=__UpperCAmelCase      )
									SCREAMING_SNAKE_CASE  :						List[str]      					= model(**__UpperCAmelCase      ).loss
									loss.backward()


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE  :						str      					= self.model_tester.prepare_config_and_inputs_for_common()

					SCREAMING_SNAKE_CASE  :						Any      					= _config_zero_init(__UpperCAmelCase      )
					for model_class in self.all_model_classes:
									SCREAMING_SNAKE_CASE  :						Optional[int]      					= model_class(config=__UpperCAmelCase      )
									for name, param in model.named_parameters():
													# we skip lambda parameters as these require special initial values
													# determined by config.layer_scale_init_value
													if "lambda" in name:
																	continue
													if param.requires_grad:
																	self.assertIn(
																	    ((param.data.mean() * 1e9).round() / 1e9).item(),					[0.0, 1.0],					msg=F'''Parameter {name} of model {model_class} seems not properly initialized''',					)






	@slow
	def 				lowercase__  (self : Dict      )     ->			Union[str, Any]:

					"""simple docstring"""

					for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									SCREAMING_SNAKE_CASE  :						Optional[Any]      					= BeitModel.from_pretrained(__UpperCAmelCase      )
									self.assertIsNotNone(__UpperCAmelCase      )






def       __lowercase  ():
				SCREAMING_SNAKE_CASE  :						int      					= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''       )
				return image




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


	@cached_property
	def 				lowercase__  (self : Any      )     ->			List[Any]:

					"""simple docstring"""

					return BeitImageProcessor.from_pretrained('''microsoft/beit-base-patch16-224'''      ) if is_vision_available() else None


	@slow
	def 				lowercase__  (self : str      )     ->			List[str]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Any      					= BeitForMaskedImageModeling.from_pretrained('''microsoft/beit-base-patch16-224-pt22k'''      ).to(__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Dict      					= self.default_image_processor
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= prepare_img()
					SCREAMING_SNAKE_CASE  :						int      					= image_processor(images=__UpperCAmelCase,					return_tensors='''pt'''      ).pixel_values.to(__UpperCAmelCase      )

					# prepare bool_masked_pos
					SCREAMING_SNAKE_CASE  :						str      					= torch.ones((1, 196),					dtype=torch.bool      ).to(__UpperCAmelCase      )

					# forward pass
					with torch.no_grad():
									SCREAMING_SNAKE_CASE  :						int      					= model(pixel_values=__UpperCAmelCase,					bool_masked_pos=__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						List[Any]      					= outputs.logits

					# verify the logits
					SCREAMING_SNAKE_CASE  :						Dict      					= torch.Size((1, 196, 8192)      )
					self.assertEqual(logits.shape,					__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= torch.tensor(
					    [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]      ).to(__UpperCAmelCase      )

					self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3],					__UpperCAmelCase,					atol=1e-2      )      )


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

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						List[Any]      					= BeitForImageClassification.from_pretrained('''microsoft/beit-base-patch16-224'''      ).to(__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						List[str]      					= self.default_image_processor
					SCREAMING_SNAKE_CASE  :						Any      					= prepare_img()
					SCREAMING_SNAKE_CASE  :						int      					= image_processor(images=__UpperCAmelCase,					return_tensors='''pt'''      ).to(__UpperCAmelCase      )

					# forward pass
					with torch.no_grad():
									SCREAMING_SNAKE_CASE  :						Any      					= model(**__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						List[str]      					= outputs.logits

					# verify the logits
					SCREAMING_SNAKE_CASE  :						List[str]      					= torch.Size((1, 1000)      )
					self.assertEqual(logits.shape,					__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= torch.tensor([-1.2385, -1.0987, -1.0108]      ).to(__UpperCAmelCase      )

					self.assertTrue(torch.allclose(logits[0, :3],					__UpperCAmelCase,					atol=1e-4      )      )

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= 281
					self.assertEqual(logits.argmax(-1      ).item(),					__UpperCAmelCase      )


	@slow
	def 				lowercase__  (self : Tuple      )     ->			str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= BeitForImageClassification.from_pretrained('''microsoft/beit-large-patch16-224-pt22k-ft22k'''      ).to(
					    __UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Any      					= self.default_image_processor
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= prepare_img()
					SCREAMING_SNAKE_CASE  :						Dict      					= image_processor(images=__UpperCAmelCase,					return_tensors='''pt'''      ).to(__UpperCAmelCase      )

					# forward pass
					with torch.no_grad():
									SCREAMING_SNAKE_CASE  :						Optional[Any]      					= model(**__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= outputs.logits

					# verify the logits
					SCREAMING_SNAKE_CASE  :						Dict      					= torch.Size((1, 21841)      )
					self.assertEqual(logits.shape,					__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Tuple      					= torch.tensor([1.6881, -0.2787, 0.5901]      ).to(__UpperCAmelCase      )

					self.assertTrue(torch.allclose(logits[0, :3],					__UpperCAmelCase,					atol=1e-4      )      )

					SCREAMING_SNAKE_CASE  :						int      					= 2396
					self.assertEqual(logits.argmax(-1      ).item(),					__UpperCAmelCase      )


	@slow
	def 				lowercase__  (self : str      )     ->			str:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Any      					= BeitForSemanticSegmentation.from_pretrained('''microsoft/beit-base-finetuned-ade-640-640'''      )
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= model.to(__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Tuple      					= BeitImageProcessor(do_resize=__UpperCAmelCase,					size=640,					do_center_crop=__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Optional[int]      					= load_dataset('''hf-internal-testing/fixtures_ade20k''',					split='''test'''      )
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= Image.open(ds[0]['''file''']      )
					SCREAMING_SNAKE_CASE  :						List[str]      					= image_processor(images=__UpperCAmelCase,					return_tensors='''pt'''      ).to(__UpperCAmelCase      )

					# forward pass
					with torch.no_grad():
									SCREAMING_SNAKE_CASE  :						Tuple      					= model(**__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						Tuple      					= outputs.logits

					# verify the logits
					SCREAMING_SNAKE_CASE  :						int      					= torch.Size((1, 150, 160, 160)      )
					self.assertEqual(logits.shape,					__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Dict      					= version.parse(PIL.__version__      ) < version.parse('''9.0.0'''      )

					if is_pillow_less_than_a:
									SCREAMING_SNAKE_CASE  :						Tuple      					= torch.tensor(
									    [
									        [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
									        [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
									        [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
									    ],					device=__UpperCAmelCase,					)
					else:
									SCREAMING_SNAKE_CASE  :						Optional[Any]      					= torch.tensor(
									    [
									        [[-4.8960, -2.3688, -3.0355], [-2.8478, -0.9836, -1.7418], [-2.9449, -1.3332, -2.1456]],
									        [[-5.8081, -3.4124, -4.1006], [-3.8561, -2.2081, -3.0323], [-3.8365, -2.4601, -3.3669]],
									        [[-0.0309, 3.9868, 4.0540], [2.9640, 4.6877, 4.9976], [3.2081, 4.7690, 4.9942]],
									    ],					device=__UpperCAmelCase,					)

					self.assertTrue(torch.allclose(logits[0, :3, :3, :3],					__UpperCAmelCase,					atol=1e-4      )      )






	@slow
	def 				lowercase__  (self : Union[str, Any]      )     ->			Union[str, Any]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						int      					= BeitForSemanticSegmentation.from_pretrained('''microsoft/beit-base-finetuned-ade-640-640'''      )
					SCREAMING_SNAKE_CASE  :						str      					= model.to(__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						int      					= BeitImageProcessor(do_resize=__UpperCAmelCase,					size=640,					do_center_crop=__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= load_dataset('''hf-internal-testing/fixtures_ade20k''',					split='''test'''      )
					SCREAMING_SNAKE_CASE  :						Any      					= Image.open(ds[0]['''file''']      )
					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= image_processor(images=__UpperCAmelCase,					return_tensors='''pt'''      ).to(__UpperCAmelCase      )

					# forward pass
					with torch.no_grad():
									SCREAMING_SNAKE_CASE  :						int      					= model(**__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Dict      					= outputs.logits.detach().cpu()

					SCREAMING_SNAKE_CASE  :						List[str]      					= image_processor.post_process_semantic_segmentation(outputs=__UpperCAmelCase,					target_sizes=[(500, 300)]      )
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= torch.Size((500, 300)      )
					self.assertEqual(segmentation[0].shape,					__UpperCAmelCase      )

					SCREAMING_SNAKE_CASE  :						Any      					= image_processor.post_process_semantic_segmentation(outputs=__UpperCAmelCase      )
					SCREAMING_SNAKE_CASE  :						str      					= torch.Size((160, 160)      )
					self.assertEqual(segmentation[0].shape,					__UpperCAmelCase      )









'''simple docstring'''

from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import ScoreSdeVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


class 	a__     (				_lowercase	):
	__magic_name__		:    UNetaDModel
	__magic_name__		:    ScoreSdeVeScheduler


	def __init__(self : List[str],					__UpperCAmelCase : UNetaDModel,					__UpperCAmelCase : ScoreSdeVeScheduler      )     ->			Tuple:

					"""simple docstring"""

					super().__init__()
					self.register_modules(unet=__UpperCAmelCase,					scheduler=__UpperCAmelCase      )






	@torch.no_grad()
	def __call__(self : int,					__UpperCAmelCase : int = 1,					__UpperCAmelCase : int = 2000,					__UpperCAmelCase : Optional[Union[torch.Generator, List[torch.Generator]]] = None,					__UpperCAmelCase : Optional[str] = "pil",					__UpperCAmelCase : bool = True,					**__UpperCAmelCase : List[str],					)     ->			Union[ImagePipelineOutput, Tuple]:

					"""simple docstring"""

					SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= self.unet.config.sample_size
					SCREAMING_SNAKE_CASE  :						Dict      					= (batch_size, 3, img_size, img_size)

					SCREAMING_SNAKE_CASE  :						List[Any]      					= self.unet

					SCREAMING_SNAKE_CASE  :						str      					= randn_tensor(__UpperCAmelCase,					generator=__UpperCAmelCase      ) * self.scheduler.init_noise_sigma
					SCREAMING_SNAKE_CASE  :						Any      					= sample.to(self.device      )

					self.scheduler.set_timesteps(__UpperCAmelCase      )
					self.scheduler.set_sigmas(__UpperCAmelCase      )

					for i, t in enumerate(self.progress_bar(self.scheduler.timesteps      )      ):
									SCREAMING_SNAKE_CASE  :						int      					= self.scheduler.sigmas[i] * torch.ones(shape[0],					device=self.device      )

									# correction step
									for _ in range(self.scheduler.config.correct_steps      ):
													SCREAMING_SNAKE_CASE  :						Any      					= self.unet(__UpperCAmelCase,					__UpperCAmelCase      ).sample
													SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= self.scheduler.step_correct(__UpperCAmelCase,					__UpperCAmelCase,					generator=__UpperCAmelCase      ).prev_sample

									# prediction step
									SCREAMING_SNAKE_CASE  :						str      					= model(__UpperCAmelCase,					__UpperCAmelCase      ).sample
									SCREAMING_SNAKE_CASE  :						Dict      					= self.scheduler.step_pred(__UpperCAmelCase,					__UpperCAmelCase,					__UpperCAmelCase,					generator=__UpperCAmelCase      )

									SCREAMING_SNAKE_CASE					,					SCREAMING_SNAKE_CASE  :						Optional[int]      					= output.prev_sample, output.prev_sample_mean

					SCREAMING_SNAKE_CASE  :						Dict      					= sample_mean.clamp(0,					1      )
					SCREAMING_SNAKE_CASE  :						str      					= sample.cpu().permute(0,					2,					3,					1      ).numpy()
					if output_type == "pil":
									SCREAMING_SNAKE_CASE  :						int      					= self.numpy_to_pil(__UpperCAmelCase      )

					if not return_dict:
									return (sample,)

					return ImagePipelineOutput(images=__UpperCAmelCase      )




