





"""simple docstring"""




import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch
		from torch import nn

		from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
		from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import ViTImageProcessor


class  __UpperCamelCase       :






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

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






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

								a   :	Union[str, Any]			=       None
								if self.use_labels:
									a   :	List[Any]			=       ids_tensor([self.batch_size]   ,      self.type_sequence_label_size   )

								a   :	Any			=       self.get_config()

								return config, pixel_values, labels






							def 					__a     (     self   )			->				List[Any]:
								return ViTConfig(
								    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=lowerCAmelCase__   ,      initializer_range=self.initializer_range   ,      encoder_stride=self.encoder_stride   ,      )






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				str:
								a   :	Union[str, Any]			=       ViTModel(config=lowerCAmelCase__   )
								model.to(lowerCAmelCase__   )
								model.eval()
								a   :	Tuple			=       model(lowerCAmelCase__   )
								self.parent.assertEqual(result.last_hidden_state.shape   ,      (self.batch_size, self.seq_length, self.hidden_size)   )






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				Union[str, Any]:
								a   :	str			=       ViTForMaskedImageModeling(config=lowerCAmelCase__   )
								model.to(lowerCAmelCase__   )
								model.eval()
								a   :	Any			=       model(lowerCAmelCase__   )
								self.parent.assertEqual(
								    result.reconstruction.shape   ,      (self.batch_size, self.num_channels, self.image_size, self.image_size)   )

								# test greyscale images
								a   :	int			=       1
								a   :	List[str]			=       ViTForMaskedImageModeling(lowerCAmelCase__   )
								model.to(lowerCAmelCase__   )
								model.eval()

								a   :	str			=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
								a   :	Any			=       model(lowerCAmelCase__   )
								self.parent.assertEqual(result.reconstruction.shape   ,      (self.batch_size, 1, self.image_size, self.image_size)   )






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				List[str]:
								a   :	Optional[int]			=       self.type_sequence_label_size
								a   :	str			=       ViTForImageClassification(lowerCAmelCase__   )
								model.to(lowerCAmelCase__   )
								model.eval()
								a   :	Optional[Any]			=       model(lowerCAmelCase__   ,      labels=lowerCAmelCase__   )
								self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.type_sequence_label_size)   )

								# test greyscale images
								a   :	int			=       1
								a   :	Dict			=       ViTForImageClassification(lowerCAmelCase__   )
								model.to(lowerCAmelCase__   )
								model.eval()

								a   :	Union[str, Any]			=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
								a   :	Dict			=       model(lowerCAmelCase__   )
								self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.type_sequence_label_size)   )






							def 					__a     (     self   )			->				Tuple:
								a   :	List[Any]			=       self.prepare_config_and_inputs()
								(
								    (
								    a
								), (
								    a
								), (
								    a
								), 
								)   :	int			=       config_and_inputs
								a   :	Any			=       {"pixel_values": pixel_values}
								return config, inputs_dict


@require_torch
class  __UpperCamelCase       (		a__	, a__	, unittest.TestCase		):
							lowerCamelCase     :      int         =(
							    (
							        ViTModel,
							        ViTForImageClassification,
							        ViTForMaskedImageModeling,
							    )
							    if is_torch_available()
							    else ()
							)
							lowerCamelCase     :      Any         =(
							    {"""feature-extraction""": ViTModel, """image-classification""": ViTForImageClassification}
							    if is_torch_available()
							    else {}
							)
							lowerCamelCase     :      Tuple         =True

							lowerCamelCase     :      Optional[Any]         =False
							lowerCamelCase     :      List[Any]         =False
							lowerCamelCase     :      Union[str, Any]         =False






							def 					__a     (     self   )			->				int:
								a   :	Optional[Any]			=       ViTModelTester(self   )
								a   :	Dict			=       ConfigTester(self   ,      config_class=lowerCAmelCase__   ,      has_text_modality=lowerCAmelCase__   ,      hidden_size=37   )






							def 					__a     (     self   )			->				List[str]:
								self.config_tester.run_common_tests()






							@unittest.skip(reason="ViT does not use inputs_embeds"   )
							def 					__a     (     self   )			->				Any:
								pass






							def 					__a     (     self   )			->				Optional[Any]:
								a, a   :	Union[str, Any]			=       self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									a   :	Optional[Any]			=       model_class(lowerCAmelCase__   )
									self.assertIsInstance(model.get_input_embeddings()   ,      (nn.Module)   )
									a   :	Any			=       model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(lowerCAmelCase__   ,      nn.Linear   )   )






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

								for model_class in self.all_model_classes:
									a   :	Optional[int]			=       model_class(lowerCAmelCase__   )
									a   :	Union[str, Any]			=       inspect.signature(model.forward   )
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									a   :	Tuple			=       [*signature.parameters.keys()]

									a   :	int			=       ["pixel_values"]
									self.assertListEqual(arg_names[:1]   ,      lowerCAmelCase__   )






							def 					__a     (     self   )			->				List[Any]:
								a   :	Optional[int]			=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCAmelCase__   )






							def 					__a     (     self   )			->				Tuple:
								a   :	List[str]			=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_image_modeling(*lowerCAmelCase__   )






							def 					__a     (     self   )			->				List[str]:
								a   :	str			=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase__   )






							@slow
							def 					__a     (     self   )			->				List[str]:
								for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									a   :	Dict			=       ViTModel.from_pretrained(lowerCAmelCase__   )
									self.assertIsNotNone(lowerCAmelCase__   )






def    _SCREAMING_SNAKE_CASE	(			)				->Any:






	'''simple docstring'''






	a   :	Optional[int]			=       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      )
	return image







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






							@cached_property
							def 					__a     (     self   )			->				str:
								return ViTImageProcessor.from_pretrained("google/vit-base-patch16-224"   ) if is_vision_available() else None






							@slow
							def 					__a     (     self   )			->				Any:
								a   :	Tuple			=       ViTForImageClassification.from_pretrained("google/vit-base-patch16-224"   ).to(lowerCAmelCase__   )

								a   :	str			=       self.default_image_processor
								a   :	Optional[Any]			=       prepare_img()
								a   :	Union[str, Any]			=       image_processor(images=lowerCAmelCase__   ,      return_tensors="pt"   ).to(lowerCAmelCase__   )

								# forward pass
								with torch.no_grad():
									a   :	str			=       model(**lowerCAmelCase__   )

								# verify the logits
								a   :	Optional[Any]			=       torch.Size((1, 1000)   )
								self.assertEqual(outputs.logits.shape   ,      lowerCAmelCase__   )

								a   :	Dict			=       torch.tensor([-0.2_744, 0.8_215, -0.0_836]   ).to(lowerCAmelCase__   )

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






							@slow
							def 					__a     (     self   )			->				Optional[Any]:
								# ViT models have an `interpolate_pos_encoding` argument in their forward method,
								# allowing to interpolate the pre-trained position embeddings in order to use
								# the model on higher resolutions. The DINO model by Facebook AI leverages this
								# to visualize self-attention on higher resolution images.
								a   :	Optional[Any]			=       ViTModel.from_pretrained("facebook/dino-vits8"   ).to(lowerCAmelCase__   )

								a   :	Union[str, Any]			=       ViTImageProcessor.from_pretrained("facebook/dino-vits8"   ,      size=480   )
								a   :	Optional[Any]			=       prepare_img()
								a   :	Any			=       image_processor(images=lowerCAmelCase__   ,      return_tensors="pt"   )
								a   :	Dict			=       inputs.pixel_values.to(lowerCAmelCase__   )

								# forward pass
								with torch.no_grad():
									a   :	Union[str, Any]			=       model(lowerCAmelCase__   ,      interpolate_pos_encoding=lowerCAmelCase__   )

								# verify the logits
								a   :	Dict			=       torch.Size((1, 3601, 384)   )
								self.assertEqual(outputs.last_hidden_state.shape   ,      lowerCAmelCase__   )

								a   :	Optional[int]			=       torch.tensor(
								    [[4.2_340, 4.3_906, -6.6_692], [4.5_463, 1.8_928, -6.7_257], [4.4_429, 0.8_496, -5.8_585]]   ).to(lowerCAmelCase__   )

								self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]   ,      lowerCAmelCase__   ,      atol=1E-4   )   )






							@slow
							@require_accelerate
							@require_torch_gpu
							def 					__a     (     self   )			->				Optional[int]:
								a   :	Tuple			=       ViTModel.from_pretrained("facebook/dino-vits8"   ,      torch_dtype=torch.floataa   ,      device_map="auto"   )
								a   :	Optional[int]			=       self.default_image_processor

								a   :	Tuple			=       prepare_img()
								a   :	List[Any]			=       image_processor(images=lowerCAmelCase__   ,      return_tensors="pt"   )
								a   :	int			=       inputs.pixel_values.to(lowerCAmelCase__   )

								# forward pass to make sure inference works in fp16
								with torch.no_grad():
									a   :	Dict			=       model(lowerCAmelCase__   )









"""simple docstring"""




from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin


@dataclass
class  __UpperCamelCase       (		a__		):
							lowerCamelCase     :      torch.FloatTensor
							lowerCamelCase     :      torch.FloatTensor
							lowerCamelCase     :      Optional[torch.FloatTensor]         =None


class  __UpperCamelCase       (		a__	, a__		):
							lowerCamelCase     :      Tuple         =2






							@register_to_config
							def __init__(     self   ,      lowerCAmelCase__ = 0.02   ,      lowerCAmelCase__ = 100   ,      lowerCAmelCase__ = 1.007   ,      lowerCAmelCase__ = 80   ,      lowerCAmelCase__ = 0.05   ,      lowerCAmelCase__ = 50   ,      )			->				Union[str, Any]:
								# standard deviation of the initial noise distribution
								a   :	Tuple			=       sigma_max

								# setable values
								a   :	int			=       None
								a   :	np.IntTensor			=       None
								a   :	torch.FloatTensor			=       None  # sigma(t_i)






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__ = None   )			->				torch.FloatTensor:
								return sample






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__ = None   )			->				List[str]:
								a   :	List[Any]			=       num_inference_steps
								a   :	List[str]			=       np.arange(0   ,      self.num_inference_steps   )[::-1].copy()
								a   :	int			=       torch.from_numpy(lowerCAmelCase__   ).to(lowerCAmelCase__   )
								a   :	List[str]			=       [
								    (
								        self.config.sigma_max**2
								        * (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1))
								    )
								    for i in self.timesteps
								]
								a   :	Any			=       torch.tensor(lowerCAmelCase__   ,      dtype=torch.floataa   ,      device=lowerCAmelCase__   )






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__ = None   )			->				Tuple[torch.FloatTensor, float]:
								if self.config.s_min <= sigma <= self.config.s_max:
									a   :	str			=       min(self.config.s_churn / self.num_inference_steps   ,      2**0.5 - 1   )
								else:
									a   :	Dict			=       0

								# sample eps ~ N(0, S_noise^2 * I)
								a   :	Union[str, Any]			=       self.config.s_noise * randn_tensor(sample.shape   ,      generator=lowerCAmelCase__   ).to(sample.device   )
								a   :	Any			=       sigma + gamma * sigma
								a   :	Tuple			=       sample + ((sigma_hat**2 - sigma**2) ** 0.5 * eps)

								return sample_hat, sigma_hat






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__ = True   ,      )			->				Union[KarrasVeOutput, Tuple]:
								a   :	Union[str, Any]			=       sample_hat + sigma_hat * model_output
								a   :	Tuple			=       (sample_hat - pred_original_sample) / sigma_hat
								a   :	List[Any]			=       sample_hat + (sigma_prev - sigma_hat) * derivative

								if not return_dict:
									return (sample_prev, derivative)

								return KarrasVeOutput(
								    prev_sample=lowerCAmelCase__   ,      derivative=lowerCAmelCase__   ,      pred_original_sample=lowerCAmelCase__   )






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__ = True   ,      )			->				Union[KarrasVeOutput, Tuple]:
								a   :	Optional[int]			=       sample_prev + sigma_prev * model_output
								a   :	str			=       (sample_prev - pred_original_sample) / sigma_prev
								a   :	Dict			=       sample_hat + (sigma_prev - sigma_hat) * (0.5 * derivative + 0.5 * derivative_corr)

								if not return_dict:
									return (sample_prev, derivative)

								return KarrasVeOutput(
								    prev_sample=lowerCAmelCase__   ,      derivative=lowerCAmelCase__   ,      pred_original_sample=lowerCAmelCase__   )






							def 					__a     (     self   ,      lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )			->				str:
								raise NotImplementedError()


