


'''simple docstring'''
import inspect
import unittest

from transformers import MobileViTVaConfig
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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation, MobileViTVaModel
						from transformers.models.mobilevitva.modeling_mobilevitva import (
						    MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
						    make_divisible,
						)


if is_vision_available():
						from PIL import Image

						from transformers import MobileViTImageProcessor






class 		UpperCamelCase_		(       A ):





						"""simple docstring"""




						def        __A       (			self					:       List[Any]	)				->       Optional[int]:
										__magic_name__							=			self.config_class(**self.inputs_dict	)
										self.parent.assertTrue(hasattr(_lowerCamelCase						,					"width_multiplier"	)	)






class 		UpperCamelCase_		:





						"""simple docstring"""




						def __init__(			self					:       Any						,					_lowerCamelCase					:       Union[str, Any]						,					_lowerCamelCase					:       List[str]=13						,					_lowerCamelCase					:       int=64						,					_lowerCamelCase					:       Any=2						,					_lowerCamelCase					:       Optional[Any]=3						,					_lowerCamelCase					:       Dict="swish"						,					_lowerCamelCase					:       str=3						,					_lowerCamelCase					:       str=32						,					_lowerCamelCase					:       Union[str, Any]=0.1						,					_lowerCamelCase					:       Tuple=0.02						,					_lowerCamelCase					:       Any=True						,					_lowerCamelCase					:       Any=True						,					_lowerCamelCase					:       List[str]=10						,					_lowerCamelCase					:       List[str]=None						,					_lowerCamelCase					:       str=0.25						,					_lowerCamelCase					:       str=0.0						,					_lowerCamelCase					:       Any=0.0						,					)				->       Dict:
										__magic_name__							=			parent
										__magic_name__							=			batch_size
										__magic_name__							=			image_size
										__magic_name__							=			patch_size
										__magic_name__							=			num_channels
										__magic_name__							=			make_divisible(5_12 * width_multiplier						,					divisor=8	)
										__magic_name__							=			hidden_act
										__magic_name__							=			conv_kernel_size
										__magic_name__							=			output_stride
										__magic_name__							=			classifier_dropout_prob
										__magic_name__							=			use_labels
										__magic_name__							=			is_training
										__magic_name__							=			num_labels
										__magic_name__							=			initializer_range
										__magic_name__							=			scope
										__magic_name__							=			width_multiplier
										__magic_name__							=			ffn_dropout
										__magic_name__							=			attn_dropout
						def        __A       (			self					:       str	)				->       Optional[int]:
										__magic_name__							=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

										__magic_name__							=			None
										__magic_name__							=			None
										if self.use_labels:
														__magic_name__							=			ids_tensor([self.batch_size]						,					self.num_labels	)
														__magic_name__							=			ids_tensor([self.batch_size, self.image_size, self.image_size]						,					self.num_labels	)

										__magic_name__							=			self.get_config()

										return config, pixel_values, labels, pixel_labels
						def        __A       (			self					:       str	)				->       Union[str, Any]:
										return MobileViTVaConfig(
										    image_size=self.image_size						,					patch_size=self.patch_size						,					num_channels=self.num_channels						,					hidden_act=self.hidden_act						,					conv_kernel_size=self.conv_kernel_size						,					output_stride=self.output_stride						,					classifier_dropout_prob=self.classifier_dropout_prob						,					initializer_range=self.initializer_range						,					width_multiplier=self.width_multiplier						,					ffn_dropout=self.ffn_dropout_prob						,					attn_dropout=self.attn_dropout_prob						,					)
						def        __A       (			self					:       Any						,					_lowerCamelCase					:       Any						,					_lowerCamelCase					:       int						,					_lowerCamelCase					:       Optional[Any]						,					_lowerCamelCase					:       str	)				->       List[Any]:
										__magic_name__							=			MobileViTVaModel(config=_lowerCamelCase	)
										model.to(_lowerCamelCase	)
										model.eval()
										__magic_name__							=			model(_lowerCamelCase	)
										self.parent.assertEqual(
										    result.last_hidden_state.shape						,					(
										        self.batch_size,
										        self.last_hidden_size,
										        self.image_size // self.output_stride,
										        self.image_size // self.output_stride,
										    )						,					)
						def        __A       (			self					:       Dict						,					_lowerCamelCase					:       int						,					_lowerCamelCase					:       Optional[int]						,					_lowerCamelCase					:       Any						,					_lowerCamelCase					:       int	)				->       str:
										__magic_name__							=			self.num_labels
										__magic_name__							=			MobileViTVaForImageClassification(_lowerCamelCase	)
										model.to(_lowerCamelCase	)
										model.eval()
										__magic_name__							=			model(_lowerCamelCase						,					labels=_lowerCamelCase	)
										self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.num_labels)	)
						def        __A       (			self					:       Optional[Any]						,					_lowerCamelCase					:       str						,					_lowerCamelCase					:       List[str]						,					_lowerCamelCase					:       Tuple						,					_lowerCamelCase					:       Dict	)				->       int:
										__magic_name__							=			self.num_labels
										__magic_name__							=			MobileViTVaForSemanticSegmentation(_lowerCamelCase	)
										model.to(_lowerCamelCase	)
										model.eval()
										__magic_name__							=			model(_lowerCamelCase	)
										self.parent.assertEqual(
										    result.logits.shape						,					(
										        self.batch_size,
										        self.num_labels,
										        self.image_size // self.output_stride,
										        self.image_size // self.output_stride,
										    )						,					)
										__magic_name__							=			model(_lowerCamelCase						,					labels=_lowerCamelCase	)
										self.parent.assertEqual(
										    result.logits.shape						,					(
										        self.batch_size,
										        self.num_labels,
										        self.image_size // self.output_stride,
										        self.image_size // self.output_stride,
										    )						,					)

						def        __A       (			self					:       List[Any]	)				->       Union[str, Any]:
										__magic_name__							=			self.prepare_config_and_inputs()
										__magic_name__							=			config_and_inputs
										__magic_name__							=			{'pixel_values': pixel_values}
										return config, inputs_dict






@require_torch
class 		UpperCamelCase_		(       A ,		A ,		unittest.TestCase ):





						"""simple docstring"""




						UpperCAmelCase__						:							Optional[int]       =  (
						    (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
						    if is_torch_available()
						    else ()
						)

						UpperCAmelCase__						:							Union[str, Any]       =  (
						    {
						        '''feature-extraction''': MobileViTVaModel,
						        '''image-classification''': MobileViTVaForImageClassification,
						        '''image-segmentation''': MobileViTVaForSemanticSegmentation,
						    }
						    if is_torch_available()
						    else {}
						)

						UpperCAmelCase__						:							Tuple       =  False
						UpperCAmelCase__						:							Any       =  False
						UpperCAmelCase__						:							str       =  False
						UpperCAmelCase__						:							Union[str, Any]       =  False
						def        __A       (			self					:       List[Any]	)				->       int:
										__magic_name__							=			MobileViTVaModelTester(self	)
										__magic_name__							=			MobileViTVaConfigTester(self						,					config_class=_lowerCamelCase						,					has_text_modality=_lowerCamelCase	)
						def        __A       (			self					:       List[str]	)				->       Optional[Any]:
										self.config_tester.run_common_tests()
						@unittest.skip(reason="MobileViTV2 does not use inputs_embeds"	)
						def        __A       (			self					:       Tuple	)				->       Union[str, Any]:
										pass
						@unittest.skip(reason="MobileViTV2 does not support input and output embeddings"	)
						def        __A       (			self					:       str	)				->       Dict:
										pass
						@unittest.skip(reason="MobileViTV2 does not output attentions"	)
						def        __A       (			self					:       Union[str, Any]	)				->       List[Any]:
										pass
						@require_torch_multi_gpu
						@unittest.skip(reason="Got `CUDA error: misaligned address` for tests after this one being run."	)
						def        __A       (			self					:       Optional[int]	)				->       Optional[Any]:
										pass
						@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."	)
						def        __A       (			self					:       Any	)				->       Optional[int]:
										pass
						def        __A       (			self					:       Dict	)				->       List[str]:
										__magic_name__							=			self.model_tester.prepare_config_and_inputs_for_common()

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

														__magic_name__							=			['pixel_values']
														self.assertListEqual(arg_names[:1]						,					_lowerCamelCase	)
						def        __A       (			self					:       List[str]	)				->       Tuple:
										__magic_name__							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_lowerCamelCase	)
						def        __A       (			self					:       Optional[int]	)				->       List[str]:
										def check_hidden_states_output(_lowerCamelCase					:       Optional[Any]						,					_lowerCamelCase					:       Union[str, Any]						,					_lowerCamelCase					:       Union[str, Any]	):
														__magic_name__							=			model_class(_lowerCamelCase	)
														model.to(_lowerCamelCase	)
														model.eval()

														with torch.no_grad():
																		__magic_name__							=			model(**self._prepare_for_class(_lowerCamelCase						,					_lowerCamelCase	)	)

														__magic_name__							=			outputs.hidden_states

														__magic_name__							=			5
														self.assertEqual(len(_lowerCamelCase	)						,					_lowerCamelCase	)

														# MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
														# with the width and height being successively divided by 2.
														__magic_name__							=			2
														for i in range(len(_lowerCamelCase	)	):
																		self.assertListEqual(
																		    list(hidden_states[i].shape[-2:]	)						,					[self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]						,					)
																		divisor *= 2

														self.assertEqual(self.model_tester.output_stride						,					divisor // 2	)

										__magic_name__							=			self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														__magic_name__							=			True
														check_hidden_states_output(_lowerCamelCase						,					_lowerCamelCase						,					_lowerCamelCase	)

														# check that output_hidden_states also work using config
														del inputs_dict["output_hidden_states"]
														__magic_name__							=			True

														check_hidden_states_output(_lowerCamelCase						,					_lowerCamelCase						,					_lowerCamelCase	)
						def        __A       (			self					:       Any	)				->       List[str]:
										__magic_name__							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase	)
						def        __A       (			self					:       List[str]	)				->       int:
										__magic_name__							=			self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_semantic_segmentation(*_lowerCamelCase	)

						@slow
						def        __A       (			self					:       Tuple	)				->       Union[str, Any]:
										for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														__magic_name__							=			MobileViTVaModel.from_pretrained(_lowerCamelCase	)
														self.assertIsNotNone(_lowerCamelCase	)





def   __snake_case							(							):
				'''simple docstring'''




				__magic_name__							=			Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
				return image






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





						"""simple docstring"""




						@cached_property
						def        __A       (			self					:       int	)				->       Tuple:
										return (
										    MobileViTImageProcessor.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256"	)
										    if is_vision_available()
										    else None
										)
						@slow
						def        __A       (			self					:       str	)				->       List[str]:
										__magic_name__							=			MobileViTVaForImageClassification.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256"	).to(
										    _lowerCamelCase	)

										__magic_name__							=			self.default_image_processor
										__magic_name__							=			prepare_img()
										__magic_name__							=			image_processor(images=_lowerCamelCase						,					return_tensors="pt"	).to(_lowerCamelCase	)

										# forward pass
										with torch.no_grad():
														__magic_name__							=			model(**_lowerCamelCase	)

										# verify the logits
										__magic_name__							=			torch.Size((1, 10_00)	)
										self.assertEqual(outputs.logits.shape						,					_lowerCamelCase	)

										__magic_name__							=			torch.tensor([-1.63_36e00, -7.32_04e-02, -5.18_83e-01]	).to(_lowerCamelCase	)

										self.assertTrue(torch.allclose(outputs.logits[0, :3]						,					_lowerCamelCase						,					atol=1e-4	)	)
						@slow
						def        __A       (			self					:       Union[str, Any]	)				->       Any:
										__magic_name__							=			MobileViTVaForSemanticSegmentation.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"	)
										__magic_name__							=			model.to(_lowerCamelCase	)

										__magic_name__							=			MobileViTImageProcessor.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"	)

										__magic_name__							=			prepare_img()
										__magic_name__							=			image_processor(images=_lowerCamelCase						,					return_tensors="pt"	).to(_lowerCamelCase	)

										# forward pass
										with torch.no_grad():
														__magic_name__							=			model(**_lowerCamelCase	)
										__magic_name__							=			outputs.logits

										# verify the logits
										__magic_name__							=			torch.Size((1, 21, 32, 32)	)
										self.assertEqual(logits.shape						,					_lowerCamelCase	)

										__magic_name__							=			torch.tensor(
										    [
										        [[7.0_863, 7.1_525, 6.8_201], [6.6_931, 6.8_770, 6.8_933], [6.2_978, 7.0_366, 6.9_636]],
										        [[-3.7_134, -3.6_712, -3.6_675], [-3.5_825, -3.3_549, -3.4_777], [-3.3_435, -3.3_979, -3.2_857]],
										        [[-2.9_329, -2.8_003, -2.7_369], [-3.0_564, -2.4_780, -2.0_207], [-2.6_889, -1.9_298, -1.7_640]],
										    ]						,					device=_lowerCamelCase						,					)

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

						@slow
						def        __A       (			self					:       List[Any]	)				->       Optional[int]:
										__magic_name__							=			MobileViTVaForSemanticSegmentation.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"	)
										__magic_name__							=			model.to(_lowerCamelCase	)

										__magic_name__							=			MobileViTImageProcessor.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"	)

										__magic_name__							=			prepare_img()
										__magic_name__							=			image_processor(images=_lowerCamelCase						,					return_tensors="pt"	).to(_lowerCamelCase	)

										# forward pass
										with torch.no_grad():
														__magic_name__							=			model(**_lowerCamelCase	)

										__magic_name__							=			outputs.logits.detach().cpu()

										__magic_name__							=			image_processor.post_process_semantic_segmentation(outputs=_lowerCamelCase						,					target_sizes=[(50, 60)]	)
										__magic_name__							=			torch.Size((50, 60)	)
										self.assertEqual(segmentation[0].shape						,					_lowerCamelCase	)

										__magic_name__							=			image_processor.post_process_semantic_segmentation(outputs=_lowerCamelCase	)
										__magic_name__							=			torch.Size((32, 32)	)
										self.assertEqual(segmentation[0].shape						,					_lowerCamelCase	)









"""simple docstring"""






import argparse

import intel_extension_for_pytorch as ipex
import torch

from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline


lowercase__			=		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""")
lowercase__			=		parser.parse_args()


lowercase__			=		"""cpu"""
lowercase__			=		"""a lovely <dicoo> in red dress and hat, in the snowly and brightly night, with many brighly buildings"""

lowercase__			=		"""path-to-your-trained-model"""
lowercase__			=		StableDiffusionPipeline.from_pretrained(model_id)
if args.dpm:
			lowercase__			=		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
lowercase__			=		pipe.to(device)

# to channels last
lowercase__			=		pipe.unet.to(memory_format=torch.channels_last)
lowercase__			=		pipe.vae.to(memory_format=torch.channels_last)
lowercase__			=		pipe.text_encoder.to(memory_format=torch.channels_last)
if pipe.requires_safety_checker:
			lowercase__			=		pipe.safety_checker.to(memory_format=torch.channels_last)

# optimize with ipex
lowercase__			=		torch.randn(2, 4, 64, 64)
lowercase__			=		torch.rand(1) * 999
lowercase__			=		torch.randn(2, 77, 768)
lowercase__			=		(sample, timestep, encoder_hidden_status)
try:
			lowercase__			=		ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True, sample_input=input_example)
except Exception:
			lowercase__			=		ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True)
lowercase__			=		ipex.optimize(pipe.vae.eval(), dtype=torch.bfloataa, inplace=True)
lowercase__			=		ipex.optimize(pipe.text_encoder.eval(), dtype=torch.bfloataa, inplace=True)
if pipe.requires_safety_checker:
			lowercase__			=		ipex.optimize(pipe.safety_checker.eval(), dtype=torch.bfloataa, inplace=True)

# compute
lowercase__			=		666
lowercase__			=		torch.Generator(device).manual_seed(seed)
lowercase__			=		{"""generator""": generator}
if args.steps is not None:
			lowercase__			=		args.steps

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

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