




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			( __UpperCAmelCase  ):





				'''simple docstring'''



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.config_class(**self.inputs_dict						)
											self.parent.assertTrue(hasattr(UpperCAmelCase_       , '''width_multiplier'''						)						)




class 							__UpperCamelCase			:





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , UpperCAmelCase_=13       , UpperCAmelCase_=64       , UpperCAmelCase_=2       , UpperCAmelCase_=3       , UpperCAmelCase_="swish"       , UpperCAmelCase_=3       , UpperCAmelCase_=32       , UpperCAmelCase_=0.1       , UpperCAmelCase_=0.02       , UpperCAmelCase_=True       , UpperCAmelCase_=True       , UpperCAmelCase_=10       , UpperCAmelCase_=None       , UpperCAmelCase_=0.25       , UpperCAmelCase_=0.0       , UpperCAmelCase_=0.0       , ):
											lowerCAmelCase   =		parent
											lowerCAmelCase   =		batch_size
											lowerCAmelCase   =		image_size
											lowerCAmelCase   =		patch_size
											lowerCAmelCase   =		num_channels
											lowerCAmelCase   =		make_divisible(5_12 * width_multiplier       , divisor=8						)
											lowerCAmelCase   =		hidden_act
											lowerCAmelCase   =		conv_kernel_size
											lowerCAmelCase   =		output_stride
											lowerCAmelCase   =		classifier_dropout_prob
											lowerCAmelCase   =		use_labels
											lowerCAmelCase   =		is_training
											lowerCAmelCase   =		num_labels
											lowerCAmelCase   =		initializer_range
											lowerCAmelCase   =		scope
											lowerCAmelCase   =		width_multiplier
											lowerCAmelCase   =		ffn_dropout
											lowerCAmelCase   =		attn_dropout



				def       __snake_case		( self						):
											lowerCAmelCase   =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

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

											lowerCAmelCase   =		self.get_config()

											return config, pixel_values, labels, pixel_labels



				def       __snake_case		( self						):
											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       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		MobileViTVaModel(config=UpperCAmelCase_						)
											model.to(UpperCAmelCase_						)
											model.eval()
											lowerCAmelCase   =		model(UpperCAmelCase_						)
											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       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		self.num_labels
											lowerCAmelCase   =		MobileViTVaForImageClassification(UpperCAmelCase_						)
											model.to(UpperCAmelCase_						)
											model.eval()
											lowerCAmelCase   =		model(UpperCAmelCase_       , labels=UpperCAmelCase_						)
											self.parent.assertEqual(result.logits.shape       , (self.batch_size, self.num_labels)						)



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											lowerCAmelCase   =		self.num_labels
											lowerCAmelCase   =		MobileViTVaForSemanticSegmentation(UpperCAmelCase_						)
											model.to(UpperCAmelCase_						)
											model.eval()
											lowerCAmelCase   =		model(UpperCAmelCase_						)
											self.parent.assertEqual(
											    result.logits.shape       , (
											        self.batch_size,
											        self.num_labels,
											        self.image_size // self.output_stride,
											        self.image_size // self.output_stride,
											    )       , )
											lowerCAmelCase   =		model(UpperCAmelCase_       , labels=UpperCAmelCase_						)
											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       __snake_case		( self						):
											lowerCAmelCase   =		self.prepare_config_and_inputs()
											lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   ,     lowerCAmelCase   =		config_and_inputs
											lowerCAmelCase   =		{'''pixel_values''': pixel_values}
											return config, inputs_dict




@require_torch
class 							__UpperCamelCase			( __UpperCAmelCase     ,					__UpperCAmelCase     ,					unittest.TestCase  ):





				'''simple docstring'''



				__a							: List[str]										=(
				    (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
				    if is_torch_available()
				    else ()
				)

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

				__a							: str										=False
				__a							: Dict										=False
				__a							: Union[str, Any]										=False
				__a							: Optional[int]										=False



				def       __snake_case		( self						):
											lowerCAmelCase   =		MobileViTVaModelTester(self						)
											lowerCAmelCase   =		MobileViTVaConfigTester(self       , config_class=UpperCAmelCase_       , has_text_modality=UpperCAmelCase_						)



				def       __snake_case		( self						):
											self.config_tester.run_common_tests()



				@unittest.skip(reason='''MobileViTV2 does not use inputs_embeds'''						)
				def       __snake_case		( self						):
											pass



				@unittest.skip(reason='''MobileViTV2 does not support input and output embeddings'''						)
				def       __snake_case		( self						):
											pass



				@unittest.skip(reason='''MobileViTV2 does not output attentions'''						)
				def       __snake_case		( self						):
											pass



				@require_torch_multi_gpu
				@unittest.skip(reason='''Got `CUDA error: misaligned address` for tests after this one being run.'''						)
				def       __snake_case		( self						):
											pass



				@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''						)
				def       __snake_case		( self						):
											pass



				def       __snake_case		( self						):
											lowerCAmelCase   ,     lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs_for_common()

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

																		lowerCAmelCase   =		['''pixel_values''']
																		self.assertListEqual(arg_names[:1]       , UpperCAmelCase_						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*UpperCAmelCase_						)



				def       __snake_case		( self						):
											def check_hidden_states_output(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
																		lowerCAmelCase   =		model_class(UpperCAmelCase_						)
																		model.to(UpperCAmelCase_						)
																		model.eval()

																		with torch.no_grad():
																									lowerCAmelCase   =		model(**self._prepare_for_class(UpperCAmelCase_       , UpperCAmelCase_						)						)

																		lowerCAmelCase   =		outputs.hidden_states

																		lowerCAmelCase   =		5
																		self.assertEqual(len(UpperCAmelCase_						)       , UpperCAmelCase_						)

																		# MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
																		# with the width and height being successively divided by 2.
																		lowerCAmelCase   =		2
																		for i in range(len(UpperCAmelCase_						)						):
																									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						)

											lowerCAmelCase   ,     lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																		lowerCAmelCase   =		True
																		check_hidden_states_output(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						)

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

																		check_hidden_states_output(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_						)



				def       __snake_case		( self						):
											lowerCAmelCase   =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_semantic_segmentation(*UpperCAmelCase_						)






				@slow
				def       __snake_case		( self						):
											for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		lowerCAmelCase   =		MobileViTVaModel.from_pretrained(UpperCAmelCase_						)
																		self.assertIsNotNone(UpperCAmelCase_						)




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


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





				'''simple docstring'''



				@cached_property
				def       __snake_case		( self						):
											return (
											    MobileViTImageProcessor.from_pretrained('''apple/mobilevitv2-1.0-imagenet1k-256'''						)
											    if is_vision_available()
											    else None
											)



				@slow
				def       __snake_case		( self						):
											lowerCAmelCase   =		MobileViTVaForImageClassification.from_pretrained('''apple/mobilevitv2-1.0-imagenet1k-256'''						).to(
											    UpperCAmelCase_						)

											lowerCAmelCase   =		self.default_image_processor
											lowerCAmelCase   =		prepare_img()
											lowerCAmelCase   =		image_processor(images=UpperCAmelCase_       , return_tensors='''pt'''						).to(UpperCAmelCase_						)

											# forward pass
											with torch.no_grad():
																		lowerCAmelCase   =		model(**UpperCAmelCase_						)

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

											lowerCAmelCase   =		torch.tensor([-1.6_3_3_6E0_0, -7.3_2_0_4E-0_2, -5.1_8_8_3E-0_1]						).to(UpperCAmelCase_						)

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



				@slow
				def       __snake_case		( self						):
											lowerCAmelCase   =		MobileViTVaForSemanticSegmentation.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''						)
											lowerCAmelCase   =		model.to(UpperCAmelCase_						)

											lowerCAmelCase   =		MobileViTImageProcessor.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''						)

											lowerCAmelCase   =		prepare_img()
											lowerCAmelCase   =		image_processor(images=UpperCAmelCase_       , return_tensors='''pt'''						).to(UpperCAmelCase_						)

											# forward pass
											with torch.no_grad():
																		lowerCAmelCase   =		model(**UpperCAmelCase_						)
											lowerCAmelCase   =		outputs.logits

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

											lowerCAmelCase   =		torch.tensor(
											    [
											        [[7.0863, 7.1525, 6.8201], [6.6931, 6.8770, 6.8933], [6.2978, 7.0366, 6.9636]],
											        [[-3.7134, -3.6712, -3.6675], [-3.5825, -3.3549, -3.4777], [-3.3435, -3.3979, -3.2857]],
											        [[-2.9329, -2.8003, -2.7369], [-3.0564, -2.4780, -2.0207], [-2.6889, -1.9298, -1.7640]],
											    ]       , device=UpperCAmelCase_       , )

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






				@slow
				def       __snake_case		( self						):
											lowerCAmelCase   =		MobileViTVaForSemanticSegmentation.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''						)
											lowerCAmelCase   =		model.to(UpperCAmelCase_						)

											lowerCAmelCase   =		MobileViTImageProcessor.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''						)

											lowerCAmelCase   =		prepare_img()
											lowerCAmelCase   =		image_processor(images=UpperCAmelCase_       , return_tensors='''pt'''						).to(UpperCAmelCase_						)

											# forward pass
											with torch.no_grad():
																		lowerCAmelCase   =		model(**UpperCAmelCase_						)

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

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

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









import torch

from diffusers import StableDiffusionPipeline


UpperCAmelCase_						="""path-to-your-trained-model"""
UpperCAmelCase_						=StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.floataa).to("""cuda""")

UpperCAmelCase_						="""A photo of sks dog in a bucket"""
UpperCAmelCase_						=pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]

image.save("""dog-bucket.png""")



