





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       lowercase					:



						"""simple docstring"""

						def __init__(     self      :		int			,					lowerCamelCase_      :		Any			,					lowerCamelCase_      :		Tuple=13			,					lowerCamelCase_      :		Any=30			,					lowerCamelCase_      :		Any=2			,					lowerCamelCase_      :		str=3			,					lowerCamelCase_      :		Any=True			,					lowerCamelCase_      :		Any=True			,					lowerCamelCase_      :		Optional[Any]=32			,					lowerCamelCase_      :		int=5			,					lowerCamelCase_      :		int=4			,					lowerCamelCase_      :		Tuple=37			,					lowerCamelCase_      :		Dict="gelu"			,					lowerCamelCase_      :		Tuple=0.1			,					lowerCamelCase_      :		Optional[Any]=0.1			,					lowerCamelCase_      :		Optional[int]=10			,					lowerCamelCase_      :		Optional[int]=0.02			,					lowerCamelCase_      :		str=None			,					lowerCamelCase_      :		int=2			,					):





										'''simple docstring'''







										_snake_case			:      Dict	    = parent
										_snake_case			:      Optional[int]	    = batch_size
										_snake_case			:      int	    = image_size
										_snake_case			:      Optional[int]	    = patch_size
										_snake_case			:      str	    = num_channels
										_snake_case			:      List[str]	    = is_training
										_snake_case			:      List[Any]	    = use_labels
										_snake_case			:      Any	    = hidden_size
										_snake_case			:      List[str]	    = num_hidden_layers
										_snake_case			:      List[str]	    = num_attention_heads
										_snake_case			:      Tuple	    = intermediate_size
										_snake_case			:      Tuple	    = hidden_act
										_snake_case			:      Optional[Any]	    = hidden_dropout_prob
										_snake_case			:      str	    = attention_probs_dropout_prob
										_snake_case			:      Any	    = type_sequence_label_size
										_snake_case			:      List[Any]	    = initializer_range
										_snake_case			:      str	    = scope
										_snake_case			:      Optional[Any]	    = encoder_stride

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







						def 		__UpperCAmelCase							(     self      :		List[str]       ):





										'''simple docstring'''







										_snake_case			:      Union[str, Any]	    = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

										_snake_case			:      Optional[int]	    = None
										if self.use_labels:
														_snake_case			:      Tuple	    = ids_tensor([self.batch_size]			,					self.type_sequence_label_size       )

										_snake_case			:      int	    = self.get_config()

										return config, pixel_values, labels







						def 		__UpperCAmelCase							(     self      :		List[str]       ):





										'''simple docstring'''







										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 		__UpperCAmelCase							(     self      :		Optional[int]			,					lowerCamelCase_      :		Union[str, Any]			,					lowerCamelCase_      :		Union[str, Any]			,					lowerCamelCase_      :		List[str]       ):





										'''simple docstring'''







										_snake_case			:      Tuple	    = ViTModel(config=lowerCamelCase_       )
										model.to(lowerCamelCase_       )
										model.eval()
										_snake_case			:      Dict	    = model(lowerCamelCase_       )
										self.parent.assertEqual(result.last_hidden_state.shape			,					(self.batch_size, self.seq_length, self.hidden_size)       )







						def 		__UpperCAmelCase							(     self      :		List[Any]			,					lowerCamelCase_      :		Optional[Any]			,					lowerCamelCase_      :		Optional[int]			,					lowerCamelCase_      :		int       ):





										'''simple docstring'''







										_snake_case			:      Optional[Any]	    = ViTForMaskedImageModeling(config=lowerCamelCase_       )
										model.to(lowerCamelCase_       )
										model.eval()
										_snake_case			:      Any	    = model(lowerCamelCase_       )
										self.parent.assertEqual(
										    result.reconstruction.shape			,					(self.batch_size, self.num_channels, self.image_size, self.image_size)       )

										# test greyscale images
										_snake_case			:      Union[str, Any]	    = 1
										_snake_case			:      Optional[int]	    = ViTForMaskedImageModeling(lowerCamelCase_       )
										model.to(lowerCamelCase_       )
										model.eval()

										_snake_case			:      Any	    = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
										_snake_case			:      List[Any]	    = model(lowerCamelCase_       )
										self.parent.assertEqual(result.reconstruction.shape			,					(self.batch_size, 1, self.image_size, self.image_size)       )







						def 		__UpperCAmelCase							(     self      :		Optional[int]			,					lowerCamelCase_      :		Tuple			,					lowerCamelCase_      :		Optional[int]			,					lowerCamelCase_      :		List[str]       ):





										'''simple docstring'''







										_snake_case			:      Optional[int]	    = self.type_sequence_label_size
										_snake_case			:      Optional[Any]	    = ViTForImageClassification(lowerCamelCase_       )
										model.to(lowerCamelCase_       )
										model.eval()
										_snake_case			:      Any	    = model(lowerCamelCase_			,					labels=lowerCamelCase_       )
										self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.type_sequence_label_size)       )

										# test greyscale images
										_snake_case			:      int	    = 1
										_snake_case			:      List[str]	    = ViTForImageClassification(lowerCamelCase_       )
										model.to(lowerCamelCase_       )
										model.eval()

										_snake_case			:      Any	    = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]       )
										_snake_case			:      Dict	    = model(lowerCamelCase_       )
										self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.type_sequence_label_size)       )







						def 		__UpperCAmelCase							(     self      :		int       ):





										'''simple docstring'''







										_snake_case			:      Dict	    = self.prepare_config_and_inputs()
										(
										    _snake_case    
										)			:      int	    = config_and_inputs
										_snake_case			:      str	    = {"pixel_values": pixel_values}
										return config, inputs_dict



@require_torch
class       lowercase					(				lowerCamelCase__		,      lowerCamelCase__		,      unittest.TestCase  ):



						"""simple docstring"""

						_UpperCamelCase	:						Optional[Any]    		=						(
						    (
						        ViTModel,
						        ViTForImageClassification,
						        ViTForMaskedImageModeling,
						    )
						    if is_torch_available()
						    else ()
						)
						_UpperCamelCase	:						List[str]    		=						(
						    {"feature-extraction": ViTModel, "image-classification": ViTForImageClassification}
						    if is_torch_available()
						    else {}
						)
						_UpperCamelCase	:						Any    		=						True

						_UpperCamelCase	:						Any    		=						False
						_UpperCamelCase	:						int    		=						False
						_UpperCamelCase	:						Optional[Any]    		=						False







						def 		__UpperCAmelCase							(     self      :		List[Any]       ):





										'''simple docstring'''







										_snake_case			:      List[Any]	    = ViTModelTester(self       )
										_snake_case			:      List[Any]	    = ConfigTester(self			,					config_class=lowerCamelCase_			,					has_text_modality=lowerCamelCase_			,					hidden_size=37       )







						def 		__UpperCAmelCase							(     self      :		int       ):





										'''simple docstring'''







										self.config_tester.run_common_tests()







						@unittest.skip(reason='ViT does not use inputs_embeds'       )
						def 		__UpperCAmelCase							(     self      :		int       ):





										'''simple docstring'''







										pass







						def 		__UpperCAmelCase							(     self      :		Tuple       ):





										'''simple docstring'''







										_snake_case			:      Optional[int]	    = self.model_tester.prepare_config_and_inputs_for_common()

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







						def 		__UpperCAmelCase							(     self      :		Optional[int]       ):





										'''simple docstring'''







										_snake_case			:      List[Any]	    = self.model_tester.prepare_config_and_inputs_for_common()

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

														_snake_case			:      Dict	    = ["pixel_values"]
														self.assertListEqual(arg_names[:1]			,					lowerCamelCase_       )







						def 		__UpperCAmelCase							(     self      :		Any       ):





										'''simple docstring'''







										_snake_case			:      int	    = self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowerCamelCase_       )







						def 		__UpperCAmelCase							(     self      :		int       ):





										'''simple docstring'''







										_snake_case			:      int	    = self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase_       )







						def 		__UpperCAmelCase							(     self      :		Union[str, Any]       ):





										'''simple docstring'''







										_snake_case			:      Dict	    = self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*lowerCamelCase_       )







						@slow
						def 		__UpperCAmelCase							(     self      :		Optional[int]       ):





										'''simple docstring'''







										for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_snake_case			:      Union[str, Any]	    = ViTModel.from_pretrained(lowerCamelCase_       )
														self.assertIsNotNone(lowerCamelCase_       )



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


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



						"""simple docstring"""

						@cached_property
						def 		__UpperCAmelCase							(     self      :		int       ):





										'''simple docstring'''







										return ViTImageProcessor.from_pretrained('google/vit-base-patch16-224'       ) if is_vision_available() else None







						@slow
						def 		__UpperCAmelCase							(     self      :		Optional[int]       ):





										'''simple docstring'''







										_snake_case			:      Optional[Any]	    = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'       ).to(lowerCamelCase_       )

										_snake_case			:      List[Any]	    = self.default_image_processor
										_snake_case			:      List[str]	    = prepare_img()
										_snake_case			:      Dict	    = image_processor(images=lowerCamelCase_			,					return_tensors='pt'       ).to(lowerCamelCase_       )

										# forward pass
										with torch.no_grad():
														_snake_case			:      str	    = model(**lowerCamelCase_       )

										# verify the logits
										_snake_case			:      Optional[Any]	    = torch.Size((1, 10_00)       )
										self.assertEqual(outputs.logits.shape			,					lowerCamelCase_       )

										_snake_case			:      Dict	    = torch.tensor([-0.2744, 0.8215, -0.0836]       ).to(lowerCamelCase_       )

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







						@slow
						def 		__UpperCAmelCase							(     self      :		Optional[Any]       ):





										'''simple docstring'''







										_snake_case			:      Optional[Any]	    = ViTModel.from_pretrained('facebook/dino-vits8'       ).to(lowerCamelCase_       )

										_snake_case			:      List[str]	    = ViTImageProcessor.from_pretrained('facebook/dino-vits8'			,					size=4_80       )
										_snake_case			:      Dict	    = prepare_img()
										_snake_case			:      Dict	    = image_processor(images=lowerCamelCase_			,					return_tensors='pt'       )
										_snake_case			:      List[Any]	    = inputs.pixel_values.to(lowerCamelCase_       )

										# forward pass
										with torch.no_grad():
														_snake_case			:      List[str]	    = model(lowerCamelCase_			,					interpolate_pos_encoding=lowerCamelCase_       )

										# verify the logits
										_snake_case			:      Tuple	    = torch.Size((1, 36_01, 3_84)       )
										self.assertEqual(outputs.last_hidden_state.shape			,					lowerCamelCase_       )

										_snake_case			:      Tuple	    = torch.tensor(
										    [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]       ).to(lowerCamelCase_       )

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







						@slow
						@require_accelerate
						@require_torch_gpu
						def 		__UpperCAmelCase							(     self      :		Optional[Any]       ):





										'''simple docstring'''







										_snake_case			:      Optional[Any]	    = ViTModel.from_pretrained('facebook/dino-vits8'			,					torch_dtype=torch.floataa			,					device_map='auto'       )
										_snake_case			:      Tuple	    = self.default_image_processor

										_snake_case			:      Optional[Any]	    = prepare_img()
										_snake_case			:      Optional[int]	    = image_processor(images=lowerCamelCase_			,					return_tensors='pt'       )
										_snake_case			:      List[Any]	    = inputs.pixel_values.to(lowerCamelCase_       )

										# forward pass to make sure inference works in fp16
										with torch.no_grad():
														_snake_case			:      Any	    = model(lowerCamelCase_       )













import argparse

import intel_extension_for_pytorch as ipex
import torch

from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline


_lowerCAmelCase				:   Any						       =						argparse.ArgumentParser('''Stable Diffusion script with intel optimization''', add_help=False)
parser.add_argument('''--dpm''', action='''store_true''', help='''Enable DPMSolver or not''')
parser.add_argument('''--steps''', default=None, type=int, help='''Num inference steps''')
_lowerCAmelCase				:   Tuple						       =						parser.parse_args()


_lowerCAmelCase				:   Any						       =						'''cpu'''
_lowerCAmelCase				:   Dict						       =						'''a lovely <dicoo> in red dress and hat, in the snowly and brightly night, with many brighly buildings'''

_lowerCAmelCase				:   Union[str, Any]						       =						'''path-to-your-trained-model'''
_lowerCAmelCase				:   Tuple						       =						StableDiffusionPipeline.from_pretrained(model_id)
if args.dpm:
 _lowerCAmelCase				:   List[str]						       =						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
_lowerCAmelCase				:   List[Any]						       =						pipe.to(device)

# to channels last
_lowerCAmelCase				:   Any						       =						pipe.unet.to(memory_format=torch.channels_last)
_lowerCAmelCase				:   int						       =						pipe.vae.to(memory_format=torch.channels_last)
_lowerCAmelCase				:   Optional[Any]						       =						pipe.text_encoder.to(memory_format=torch.channels_last)
if pipe.requires_safety_checker:
 _lowerCAmelCase				:   Optional[int]						       =						pipe.safety_checker.to(memory_format=torch.channels_last)

# optimize with ipex
_lowerCAmelCase				:   Tuple						       =						torch.randn(2, 4, 64, 64)
_lowerCAmelCase				:   List[str]						       =						torch.rand(1) * 999
_lowerCAmelCase				:   List[Any]						       =						torch.randn(2, 77, 768)
_lowerCAmelCase				:   int						       =						(sample, timestep, encoder_hidden_status)
try:
 _lowerCAmelCase				:   Tuple						       =						ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True, sample_input=input_example)
except Exception:
 _lowerCAmelCase				:   Optional[Any]						       =						ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True)
_lowerCAmelCase				:   Dict						       =						ipex.optimize(pipe.vae.eval(), dtype=torch.bfloataa, inplace=True)
_lowerCAmelCase				:   Dict						       =						ipex.optimize(pipe.text_encoder.eval(), dtype=torch.bfloataa, inplace=True)
if pipe.requires_safety_checker:
 _lowerCAmelCase				:   Tuple						       =						ipex.optimize(pipe.safety_checker.eval(), dtype=torch.bfloataa, inplace=True)

# compute
_lowerCAmelCase				:   Dict						       =						666
_lowerCAmelCase				:   Tuple						       =						torch.Generator(device).manual_seed(seed)
_lowerCAmelCase				:   Tuple						       =						{'''generator''': generator}
if args.steps is not None:
 _lowerCAmelCase				:   Tuple						       =						args.steps

with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloataa):
 _lowerCAmelCase				:   int						       =						pipe(prompt, **generate_kwargs).images[0]

# save image
image.save('''generated.png''')
