




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class       UpperCAmelCase_   ( A		,   unittest.TestCase							):




							'''simple docstring'''



							a__    	=  UnCLIPPipeline
							a__    	=  TEXT_TO_IMAGE_PARAMS - {
							    '''negative_prompt''',
							    '''height''',
							    '''width''',
							    '''negative_prompt_embeds''',
							    '''guidance_scale''',
							    '''prompt_embeds''',
							    '''cross_attention_kwargs''',
							}
							a__    	=  TEXT_TO_IMAGE_BATCH_PARAMS
							a__    	=  [
							    '''generator''',
							    '''return_dict''',
							    '''prior_num_inference_steps''',
							    '''decoder_num_inference_steps''',
							    '''super_res_num_inference_steps''',
							]
							a__    	=  False





							@property
							def  _UpperCAmelCase  (       self							:						Any      )      ->       List[Any]:
										return 32





							@property
							def  _UpperCAmelCase  (       self							:						str      )      ->       Union[str, Any]:
										return 32





							@property
							def  _UpperCAmelCase  (       self							:						str      )      ->       List[Any]:
										return self.time_input_dim





							@property
							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       Union[str, Any]:
										return self.time_input_dim * 4





							@property
							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       Dict:
										return 100





							@property
							def  _UpperCAmelCase  (       self							:						List[Any]      )      ->       Dict:
										SCREAMING_SNAKE_CASE     =       CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )
										return tokenizer





							@property
							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       Optional[int]:
										torch.manual_seed(0      )
										SCREAMING_SNAKE_CASE     =       CLIPTextConfig(
										    bos_token_id=0							,  eos_token_id=2							,  hidden_size=self.text_embedder_hidden_size							,  projection_dim=self.text_embedder_hidden_size							,  intermediate_size=37							,  layer_norm_eps=1E-0_5							,  num_attention_heads=4							,  num_hidden_layers=5							,  pad_token_id=1							,  vocab_size=1_000							,  )
										return CLIPTextModelWithProjection(a      )





							@property
							def  _UpperCAmelCase  (       self							:						Dict      )      ->       Dict:
										torch.manual_seed(0      )

										SCREAMING_SNAKE_CASE     =       {
										    """num_attention_heads""": 2,
										    """attention_head_dim""": 12,
										    """embedding_dim""": self.text_embedder_hidden_size,
										    """num_layers""": 1,
										}

										SCREAMING_SNAKE_CASE     =       PriorTransformer(**a      )
										return model





							@property
							def  _UpperCAmelCase  (       self							:						Tuple      )      ->       List[Any]:
										torch.manual_seed(0      )

										SCREAMING_SNAKE_CASE     =       {
										    """clip_embeddings_dim""": self.text_embedder_hidden_size,
										    """time_embed_dim""": self.time_embed_dim,
										    """cross_attention_dim""": self.cross_attention_dim,
										}

										SCREAMING_SNAKE_CASE     =       UnCLIPTextProjModel(**a      )
										return model





							@property
							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       List[Any]:
										torch.manual_seed(0      )

										SCREAMING_SNAKE_CASE     =       {
										    """sample_size""": 32,
										    # RGB in channels
										    """in_channels""": 3,
										    # Out channels is double in channels because predicts mean and variance
										    """out_channels""": 6,
										    """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
										    """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
										    """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
										    """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
										    """layers_per_block""": 1,
										    """cross_attention_dim""": self.cross_attention_dim,
										    """attention_head_dim""": 4,
										    """resnet_time_scale_shift""": """scale_shift""",
										    """class_embed_type""": """identity""",
										}

										SCREAMING_SNAKE_CASE     =       UNetaDConditionModel(**a      )
										return model





							@property
							def  _UpperCAmelCase  (       self							:						Any      )      ->       List[Any]:
										return {
										    "sample_size": 64,
										    "layers_per_block": 1,
										    "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
										    "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
										    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
										    "in_channels": 6,
										    "out_channels": 3,
										}





							@property
							def  _UpperCAmelCase  (       self							:						Union[str, Any]      )      ->       int:
										torch.manual_seed(0      )

										SCREAMING_SNAKE_CASE     =       UNetaDModel(**self.dummy_super_res_kwargs      )
										return model





							@property
							def  _UpperCAmelCase  (       self							:						Optional[Any]      )      ->       Optional[int]:
										# seeded differently to get different unet than `self.dummy_super_res_first`
										torch.manual_seed(1      )

										SCREAMING_SNAKE_CASE     =       UNetaDModel(**self.dummy_super_res_kwargs      )
										return model





							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       Optional[Any]:
										SCREAMING_SNAKE_CASE     =       self.dummy_prior
										SCREAMING_SNAKE_CASE     =       self.dummy_decoder
										SCREAMING_SNAKE_CASE     =       self.dummy_text_proj
										SCREAMING_SNAKE_CASE     =       self.dummy_text_encoder
										SCREAMING_SNAKE_CASE     =       self.dummy_tokenizer
										SCREAMING_SNAKE_CASE     =       self.dummy_super_res_first
										SCREAMING_SNAKE_CASE     =       self.dummy_super_res_last

										SCREAMING_SNAKE_CASE     =       UnCLIPScheduler(
										    variance_type="""fixed_small_log"""							,  prediction_type="""sample"""							,  num_train_timesteps=1_000							,  clip_sample_range=5.0							,  )

										SCREAMING_SNAKE_CASE     =       UnCLIPScheduler(
										    variance_type="""learned_range"""							,  prediction_type="""epsilon"""							,  num_train_timesteps=1_000							,  )

										SCREAMING_SNAKE_CASE     =       UnCLIPScheduler(
										    variance_type="""fixed_small_log"""							,  prediction_type="""epsilon"""							,  num_train_timesteps=1_000							,  )

										SCREAMING_SNAKE_CASE     =       {
										    """prior""": prior,
										    """decoder""": decoder,
										    """text_proj""": text_proj,
										    """text_encoder""": text_encoder,
										    """tokenizer""": tokenizer,
										    """super_res_first""": super_res_first,
										    """super_res_last""": super_res_last,
										    """prior_scheduler""": prior_scheduler,
										    """decoder_scheduler""": decoder_scheduler,
										    """super_res_scheduler""": super_res_scheduler,
										}

										return components





							def  _UpperCAmelCase  (       self							:						List[str]							,  a							:						List[str]							,  a							:						Union[str, Any]=0      )      ->       Tuple:
										if str(a      ).startswith("""mps"""      ):
													SCREAMING_SNAKE_CASE     =       torch.manual_seed(a      )
										else:
													SCREAMING_SNAKE_CASE     =       torch.Generator(device=a      ).manual_seed(a      )
										SCREAMING_SNAKE_CASE     =       {
										    """prompt""": """horse""",
										    """generator""": generator,
										    """prior_num_inference_steps""": 2,
										    """decoder_num_inference_steps""": 2,
										    """super_res_num_inference_steps""": 2,
										    """output_type""": """numpy""",
										}
										return inputs





							def  _UpperCAmelCase  (       self							:						Any      )      ->       str:
										SCREAMING_SNAKE_CASE     =       """cpu"""

										SCREAMING_SNAKE_CASE     =       self.get_dummy_components()

										SCREAMING_SNAKE_CASE     =       self.pipeline_class(**a      )
										SCREAMING_SNAKE_CASE     =       pipe.to(a      )

										pipe.set_progress_bar_config(disable=a      )

										SCREAMING_SNAKE_CASE     =       pipe(**self.get_dummy_inputs(a      )      )
										SCREAMING_SNAKE_CASE     =       output.images

										SCREAMING_SNAKE_CASE     =       pipe(
										    **self.get_dummy_inputs(a      )							,  return_dict=a							,  )[0]

										SCREAMING_SNAKE_CASE     =       image[0, -3:, -3:, -1]
										SCREAMING_SNAKE_CASE     =       image_from_tuple[0, -3:, -3:, -1]

										assert image.shape == (1, 64, 64, 3)

										SCREAMING_SNAKE_CASE     =       np.array(
										    [
										        0.9997,
										        0.9988,
										        0.0028,
										        0.9997,
										        0.9984,
										        0.9965,
										        0.0029,
										        0.9986,
										        0.0025,
										    ]      )

										assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2
										assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 1E-2





							def  _UpperCAmelCase  (       self							:						Optional[Any]      )      ->       Dict:
										SCREAMING_SNAKE_CASE     =       torch.device("""cpu"""      )




										class       UpperCAmelCase_   :




																	'''simple docstring'''



																	a__    	=  1





										SCREAMING_SNAKE_CASE     =       self.get_dummy_components()

										SCREAMING_SNAKE_CASE     =       self.pipeline_class(**a      )
										SCREAMING_SNAKE_CASE     =       pipe.to(a      )

										SCREAMING_SNAKE_CASE     =       components["""prior"""]
										SCREAMING_SNAKE_CASE     =       components["""decoder"""]
										SCREAMING_SNAKE_CASE     =       components["""super_res_first"""]
										SCREAMING_SNAKE_CASE     =       components["""tokenizer"""]
										SCREAMING_SNAKE_CASE     =       components["""text_encoder"""]

										SCREAMING_SNAKE_CASE     =       torch.Generator(device=a      ).manual_seed(0      )
										SCREAMING_SNAKE_CASE     =       prior.dtype
										SCREAMING_SNAKE_CASE     =       1

										SCREAMING_SNAKE_CASE     =       (batch_size, prior.config.embedding_dim)
										SCREAMING_SNAKE_CASE     =       pipe.prepare_latents(
										    a							,  dtype=a							,  device=a							,  generator=a							,  latents=a							,  scheduler=DummyScheduler()      )
										SCREAMING_SNAKE_CASE     =       (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
										SCREAMING_SNAKE_CASE     =       pipe.prepare_latents(
										    a							,  dtype=a							,  device=a							,  generator=a							,  latents=a							,  scheduler=DummyScheduler()      )

										SCREAMING_SNAKE_CASE     =       (
										    batch_size,
										    super_res_first.config.in_channels // 2,
										    super_res_first.config.sample_size,
										    super_res_first.config.sample_size,
										)
										SCREAMING_SNAKE_CASE     =       pipe.prepare_latents(
										    a							,  dtype=a							,  device=a							,  generator=a							,  latents=a							,  scheduler=DummyScheduler()      )

										pipe.set_progress_bar_config(disable=a      )

										SCREAMING_SNAKE_CASE     =       """this is a prompt example"""

										SCREAMING_SNAKE_CASE     =       torch.Generator(device=a      ).manual_seed(0      )
										SCREAMING_SNAKE_CASE     =       pipe(
										    [prompt]							,  generator=a							,  prior_num_inference_steps=2							,  decoder_num_inference_steps=2							,  super_res_num_inference_steps=2							,  prior_latents=a							,  decoder_latents=a							,  super_res_latents=a							,  output_type="""np"""							,  )
										SCREAMING_SNAKE_CASE     =       output.images

										SCREAMING_SNAKE_CASE     =       tokenizer(
										    a							,  padding="""max_length"""							,  max_length=tokenizer.model_max_length							,  return_tensors="""pt"""							,  )
										SCREAMING_SNAKE_CASE     =       text_encoder(text_inputs.input_ids      )
										SCREAMING_SNAKE_CASE     =       text_inputs.attention_mask

										SCREAMING_SNAKE_CASE     =       torch.Generator(device=a      ).manual_seed(0      )
										SCREAMING_SNAKE_CASE     =       pipe(
										    generator=a							,  prior_num_inference_steps=2							,  decoder_num_inference_steps=2							,  super_res_num_inference_steps=2							,  prior_latents=a							,  decoder_latents=a							,  super_res_latents=a							,  text_model_output=a							,  text_attention_mask=a							,  output_type="""np"""							,  )[0]

										# make sure passing text embeddings manually is identical
										assert np.abs(image - image_from_text      ).max() < 1E-4





							@skip_mps
							def  _UpperCAmelCase  (       self							:						Dict      )      ->       Optional[int]:
										SCREAMING_SNAKE_CASE     =       torch_device == """cpu"""

										self._test_attention_slicing_forward_pass(test_max_difference=a							,  expected_max_diff=0.01      )





							@skip_mps
							def  _UpperCAmelCase  (       self							:						Union[str, Any]      )      ->       Union[str, Any]:
										SCREAMING_SNAKE_CASE     =       torch_device == """cpu"""
										SCREAMING_SNAKE_CASE     =       True
										SCREAMING_SNAKE_CASE     =       [
										    """prior_num_inference_steps""",
										    """decoder_num_inference_steps""",
										    """super_res_num_inference_steps""",
										]

										self._test_inference_batch_single_identical(
										    test_max_difference=a							,  relax_max_difference=a							,  additional_params_copy_to_batched_inputs=a							,  )





							def  _UpperCAmelCase  (       self							:						int      )      ->       List[Any]:
										SCREAMING_SNAKE_CASE     =       [
										    """prior_num_inference_steps""",
										    """decoder_num_inference_steps""",
										    """super_res_num_inference_steps""",
										]

										if torch_device == "mps":
													# TODO: MPS errors with larger batch sizes
													SCREAMING_SNAKE_CASE     =       [2, 3]
													self._test_inference_batch_consistent(
													    batch_sizes=a							,  additional_params_copy_to_batched_inputs=a							,  )
										else:
													self._test_inference_batch_consistent(
													    additional_params_copy_to_batched_inputs=a      )





							@skip_mps
							def  _UpperCAmelCase  (       self							:						Optional[Any]      )      ->       Optional[Any]:
										return super().test_dict_tuple_outputs_equivalent()





							@skip_mps
							def  _UpperCAmelCase  (       self							:						Union[str, Any]      )      ->       List[Any]:
										return super().test_save_load_local()







							@skip_mps
							def  _UpperCAmelCase  (       self							:						Tuple      )      ->       Dict:
										return super().test_save_load_optional_components()





@nightly
class       UpperCAmelCase_   ( unittest.TestCase							):




							'''simple docstring'''



							def  _UpperCAmelCase  (       self							:						Dict      )      ->       Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()







							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       int:
										SCREAMING_SNAKE_CASE     =       load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
										    """/unclip/karlo_v1_alpha_horse_cpu.npy"""      )

										SCREAMING_SNAKE_CASE     =       UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""      )
										pipeline.set_progress_bar_config(disable=a      )

										SCREAMING_SNAKE_CASE     =       torch.manual_seed(0      )
										SCREAMING_SNAKE_CASE     =       pipeline(
										    """horse"""							,  num_images_per_prompt=1							,  generator=a							,  output_type="""np"""							,  )

										SCREAMING_SNAKE_CASE     =       output.images[0]

										assert image.shape == (256, 256, 3)
										assert np.abs(expected_image - image      ).max() < 1E-1





@slow
@require_torch_gpu
class       UpperCAmelCase_   ( unittest.TestCase							):




							'''simple docstring'''



							def  _UpperCAmelCase  (       self							:						List[Any]      )      ->       Optional[Any]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()





							def  _UpperCAmelCase  (       self							:						List[Any]      )      ->       List[Any]:
										SCREAMING_SNAKE_CASE     =       load_numpy(
										    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
										    """/unclip/karlo_v1_alpha_horse_fp16.npy"""      )

										SCREAMING_SNAKE_CASE     =       UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""							,  torch_dtype=torch.floataa      )
										SCREAMING_SNAKE_CASE     =       pipeline.to(a      )
										pipeline.set_progress_bar_config(disable=a      )

										SCREAMING_SNAKE_CASE     =       torch.Generator(device="""cpu"""      ).manual_seed(0      )
										SCREAMING_SNAKE_CASE     =       pipeline(
										    """horse"""							,  generator=a							,  output_type="""np"""							,  )

										SCREAMING_SNAKE_CASE     =       output.images[0]

										assert image.shape == (256, 256, 3)

										assert_mean_pixel_difference(a							,  a      )







							def  _UpperCAmelCase  (       self							:						str      )      ->       Dict:
										torch.cuda.empty_cache()
										torch.cuda.reset_max_memory_allocated()
										torch.cuda.reset_peak_memory_stats()

										SCREAMING_SNAKE_CASE     =       UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""							,  torch_dtype=torch.floataa      )
										SCREAMING_SNAKE_CASE     =       pipe.to(a      )
										pipe.set_progress_bar_config(disable=a      )
										pipe.enable_attention_slicing()
										pipe.enable_sequential_cpu_offload()

										SCREAMING_SNAKE_CASE     =       pipe(
										    """horse"""							,  num_images_per_prompt=1							,  prior_num_inference_steps=2							,  decoder_num_inference_steps=2							,  super_res_num_inference_steps=2							,  output_type="""np"""							,  )

										SCREAMING_SNAKE_CASE     =       torch.cuda.max_memory_allocated()
										# make sure that less than 7 GB is allocated
										assert mem_bytes < 7 * 10**9












from __future__ import annotations

import inspect
import unittest

from transformers import MobileViTConfig
from transformers.file_utils import is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import numpy as np
		import tensorflow as tf

		from transformers import TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation, TFMobileViTModel
		from transformers.models.mobilevit.modeling_tf_mobilevit import TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import MobileViTImageProcessor




class       UpperCAmelCase_   ( A							):




							'''simple docstring'''



							def  _UpperCAmelCase  (       self							:						Tuple      )      ->       int:
										SCREAMING_SNAKE_CASE     =       self.config_class(**self.inputs_dict      )
										self.parent.assertTrue(hasattr(a							,  """hidden_sizes"""      )      )
										self.parent.assertTrue(hasattr(a							,  """neck_hidden_sizes"""      )      )
										self.parent.assertTrue(hasattr(a							,  """num_attention_heads"""      )      )





class       UpperCAmelCase_   :




							'''simple docstring'''



							def __init__(       self							:						Optional[Any]							,  a							:						Dict							,  a							:						Optional[int]=13							,  a							:						Optional[int]=32							,  a							:						List[str]=2							,  a							:						Optional[int]=3							,  a							:						str=640							,  a							:						str=4							,  a							:						Optional[int]="silu"							,  a							:						Union[str, Any]=3							,  a							:						Optional[int]=32							,  a							:						Optional[Any]=0.1							,  a							:						Optional[int]=0.1							,  a							:						int=0.1							,  a							:						List[str]=0.02							,  a							:						Tuple=True							,  a							:						Optional[Any]=True							,  a							:						List[Any]=10							,  a							:						List[Any]=None							,  )      ->       List[str]:
										SCREAMING_SNAKE_CASE     =       parent
										SCREAMING_SNAKE_CASE     =       batch_size
										SCREAMING_SNAKE_CASE     =       image_size
										SCREAMING_SNAKE_CASE     =       patch_size
										SCREAMING_SNAKE_CASE     =       num_channels
										SCREAMING_SNAKE_CASE     =       last_hidden_size
										SCREAMING_SNAKE_CASE     =       num_attention_heads
										SCREAMING_SNAKE_CASE     =       hidden_act
										SCREAMING_SNAKE_CASE     =       conv_kernel_size
										SCREAMING_SNAKE_CASE     =       output_stride
										SCREAMING_SNAKE_CASE     =       hidden_dropout_prob
										SCREAMING_SNAKE_CASE     =       attention_probs_dropout_prob
										SCREAMING_SNAKE_CASE     =       classifier_dropout_prob
										SCREAMING_SNAKE_CASE     =       use_labels
										SCREAMING_SNAKE_CASE     =       is_training
										SCREAMING_SNAKE_CASE     =       num_labels
										SCREAMING_SNAKE_CASE     =       initializer_range
										SCREAMING_SNAKE_CASE     =       scope





							def  _UpperCAmelCase  (       self							:						Tuple      )      ->       List[Any]:
										SCREAMING_SNAKE_CASE     =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

										SCREAMING_SNAKE_CASE     =       self.get_config()

										return config, pixel_values, labels, pixel_labels





							def  _UpperCAmelCase  (       self							:						List[Any]      )      ->       Union[str, Any]:
										return MobileViTConfig(
										    image_size=self.image_size							,  patch_size=self.patch_size							,  num_channels=self.num_channels							,  num_attention_heads=self.num_attention_heads							,  hidden_act=self.hidden_act							,  conv_kernel_size=self.conv_kernel_size							,  output_stride=self.output_stride							,  hidden_dropout_prob=self.hidden_dropout_prob							,  attention_probs_dropout_prob=self.attention_probs_dropout_prob							,  classifier_dropout_prob=self.classifier_dropout_prob							,  initializer_range=self.initializer_range							,  )





							def  _UpperCAmelCase  (       self							:						List[str]							,  a							:						Union[str, Any]							,  a							:						Tuple							,  a							:						Optional[int]							,  a							:						str      )      ->       Tuple:
										SCREAMING_SNAKE_CASE     =       TFMobileViTModel(config=a      )
										SCREAMING_SNAKE_CASE     =       model(a							,  training=a      )
										SCREAMING_SNAKE_CASE     =       SCREAMING_SNAKE_CASE     =       self.image_size // self.output_stride
										self.parent.assertEqual(
										    result.last_hidden_state.shape							,  (self.batch_size, self.last_hidden_size, expected_height, expected_width)      )





							def  _UpperCAmelCase  (       self							:						Any							,  a							:						Dict							,  a							:						Optional[Any]							,  a							:						List[Any]							,  a							:						List[str]      )      ->       Any:
										SCREAMING_SNAKE_CASE     =       self.num_labels
										SCREAMING_SNAKE_CASE     =       TFMobileViTForImageClassification(a      )
										SCREAMING_SNAKE_CASE     =       model(a							,  labels=a							,  training=a      )
										self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.num_labels)      )





							def  _UpperCAmelCase  (       self							:						Dict							,  a							:						List[str]							,  a							:						Union[str, Any]							,  a							:						List[Any]							,  a							:						Union[str, Any]      )      ->       Optional[int]:
										SCREAMING_SNAKE_CASE     =       self.num_labels
										SCREAMING_SNAKE_CASE     =       TFMobileViTForSemanticSegmentation(a      )
										SCREAMING_SNAKE_CASE     =       SCREAMING_SNAKE_CASE     =       self.image_size // self.output_stride

										SCREAMING_SNAKE_CASE     =       model(a							,  training=a      )
										self.parent.assertEqual(
										    result.logits.shape							,  (self.batch_size, self.num_labels, expected_height, expected_width)      )

										SCREAMING_SNAKE_CASE     =       model(a							,  labels=a							,  training=a      )
										self.parent.assertEqual(
										    result.logits.shape							,  (self.batch_size, self.num_labels, expected_height, expected_width)      )







							def  _UpperCAmelCase  (       self							:						Any      )      ->       Optional[int]:
										SCREAMING_SNAKE_CASE     =       self.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE     =       config_and_inputs
										SCREAMING_SNAKE_CASE     =       {"""pixel_values""": pixel_values}
										return config, inputs_dict





@require_tf
class       UpperCAmelCase_   ( A		,   A		,   unittest.TestCase							):




							'''simple docstring'''



							a__    	=  (
							    (TFMobileViTModel, TFMobileViTForImageClassification, TFMobileViTForSemanticSegmentation)
							    if is_tf_available()
							    else ()
							)
							a__    	=  (
							    {'''feature-extraction''': TFMobileViTModel, '''image-classification''': TFMobileViTForImageClassification}
							    if is_tf_available()
							    else {}
							)

							a__    	=  False
							a__    	=  False
							a__    	=  False
							a__    	=  False
							a__    	=  False





							def  _UpperCAmelCase  (       self							:						str      )      ->       Dict:
										SCREAMING_SNAKE_CASE     =       TFMobileViTModelTester(self      )
										SCREAMING_SNAKE_CASE     =       TFMobileViTConfigTester(self							,  config_class=a							,  has_text_modality=a      )





							def  _UpperCAmelCase  (       self							:						Optional[Any]      )      ->       Dict:
										self.config_tester.run_common_tests()





							@unittest.skip(reason="""MobileViT does not use inputs_embeds"""      )
							def  _UpperCAmelCase  (       self							:						Dict      )      ->       Tuple:
										pass





							@unittest.skip(reason="""MobileViT does not support input and output embeddings"""      )
							def  _UpperCAmelCase  (       self							:						List[Any]      )      ->       Tuple:
										pass





							@unittest.skip(reason="""MobileViT does not output attentions"""      )
							def  _UpperCAmelCase  (       self							:						str      )      ->       List[Any]:
										pass





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

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

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





							def  _UpperCAmelCase  (       self							:						Union[str, Any]      )      ->       List[Any]:
										SCREAMING_SNAKE_CASE     =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*a      )





							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       Union[str, Any]:
										def check_hidden_states_output(a							:						Dict							,  a							:						Union[str, Any]							,  a							:						Optional[Any]      ):
													SCREAMING_SNAKE_CASE     =       model_class(a      )

													SCREAMING_SNAKE_CASE     =       model(**self._prepare_for_class(a							,  a      )      )

													SCREAMING_SNAKE_CASE     =       outputs.hidden_states

													SCREAMING_SNAKE_CASE     =       5
													self.assertEqual(len(a      )							,  a      )

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

										SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE     =       self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													SCREAMING_SNAKE_CASE     =       True
													check_hidden_states_output(a							,  a							,  a      )

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

													check_hidden_states_output(a							,  a							,  a      )





							def  _UpperCAmelCase  (       self							:						Optional[Any]      )      ->       List[str]:
										SCREAMING_SNAKE_CASE     =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*a      )





							def  _UpperCAmelCase  (       self							:						List[str]      )      ->       Tuple:
										SCREAMING_SNAKE_CASE     =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_semantic_segmentation(*a      )





							@unittest.skipIf(
							    not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""      )      ) == 0							,  reason="""TF does not support backprop for grouped convolutions on CPU."""							,  )
							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       Any:
										super().test_dataset_conversion()





							def  _UpperCAmelCase  (       self							:						Union[str, Any]							,  a							:						List[str]							,  a							:						Optional[int]							,  a							:						Any=2E-1							,  a							:						List[str]=2E-1      )      ->       Dict:
										self.assertTrue(np.allclose(a							,  a							,  atol=a							,  rtol=a      )      )





							@unittest.skipIf(
							    not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""      )      ) == 0							,  reason="""TF does not support backprop for grouped convolutions on CPU."""							,  )
							@slow
							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       str:
										SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE     =       self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													# Since `TFMobileViTModel` cannot operate with the default `fit()` method.
													if model_class.__name__ != "TFMobileViTModel":
																SCREAMING_SNAKE_CASE     =       model_class(a      )
																if getattr(a							,  """hf_compute_loss"""							,  a      ):
																			super().test_keras_fit()





							def  _UpperCAmelCase  (       self							:						int      )      ->       Optional[Any]:
										SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE     =       self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes:
													# set an ignore index to correctly test the masked loss used in
													# `TFMobileViTForSemanticSegmentation`.
													if model_class.__name__ != "TFMobileViTForSemanticSegmentation":
																SCREAMING_SNAKE_CASE     =       5

													SCREAMING_SNAKE_CASE     =       model_class(a      )
													if getattr(a							,  """hf_compute_loss"""							,  a      ):
																# The number of elements in the loss should be the same as the number of elements in the label
																SCREAMING_SNAKE_CASE     =       self._prepare_for_class(inputs_dict.copy()							,  a							,  return_labels=a      )
																SCREAMING_SNAKE_CASE     =       prepared_for_class[
																    sorted(prepared_for_class.keys() - inputs_dict.keys()							,  reverse=a      )[0]
																]
																SCREAMING_SNAKE_CASE     =       added_label.shape.as_list()[:1]

																# Test that model correctly compute the loss with kwargs
																SCREAMING_SNAKE_CASE     =       self._prepare_for_class(inputs_dict.copy()							,  a							,  return_labels=a      )
																SCREAMING_SNAKE_CASE     =       {"""input_ids""", """pixel_values""", """input_features"""}
																SCREAMING_SNAKE_CASE     =       possible_input_names.intersection(set(a      )      ).pop()
																SCREAMING_SNAKE_CASE     =       prepared_for_class.pop(a      )

																SCREAMING_SNAKE_CASE     =       model(a							,  **a      )[0]
																self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]      )

																# Test that model correctly compute the loss when we mask some positions
																SCREAMING_SNAKE_CASE     =       self._prepare_for_class(inputs_dict.copy()							,  a							,  return_labels=a      )
																SCREAMING_SNAKE_CASE     =       {"""input_ids""", """pixel_values""", """input_features"""}
																SCREAMING_SNAKE_CASE     =       possible_input_names.intersection(set(a      )      ).pop()
																SCREAMING_SNAKE_CASE     =       prepared_for_class.pop(a      )
																if "labels" in prepared_for_class:
																			SCREAMING_SNAKE_CASE     =       prepared_for_class["""labels"""].numpy()
																			if len(labels.shape      ) > 1 and labels.shape[1] != 1:
																						# labels[0] = -100
																						SCREAMING_SNAKE_CASE     =       tf.convert_to_tensor(a      )
																						SCREAMING_SNAKE_CASE     =       model(a							,  **a      )[0]
																						self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]      )
																						self.assertTrue(not np.any(np.isnan(loss.numpy()      )      )      )

                # Test that model correctly compute the loss with a dict
																SCREAMING_SNAKE_CASE     =       self._prepare_for_class(inputs_dict.copy()							,  a							,  return_labels=a      )
																SCREAMING_SNAKE_CASE     =       model(a      )[0]
																self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]      )

																# Test that model correctly compute the loss with a tuple
																SCREAMING_SNAKE_CASE     =       self._prepare_for_class(inputs_dict.copy()							,  a							,  return_labels=a      )

																# Get keys that were added with the _prepare_for_class function
																SCREAMING_SNAKE_CASE     =       prepared_for_class.keys() - inputs_dict.keys()
																SCREAMING_SNAKE_CASE     =       inspect.signature(model.call      ).parameters
																SCREAMING_SNAKE_CASE     =       list(signature.keys()      )

																# Create a dictionary holding the location of the tensors in the tuple
																SCREAMING_SNAKE_CASE     =       {0: input_name}
																for label_key in label_keys:
																			SCREAMING_SNAKE_CASE     =       signature_names.index(a      )
																			SCREAMING_SNAKE_CASE     =       label_key
																SCREAMING_SNAKE_CASE     =       sorted(tuple_index_mapping.items()      )
																# Initialize a list with their default values, update the values and convert to a tuple
																SCREAMING_SNAKE_CASE     =       []

																for name in signature_names:
																			if name != "kwargs":
																						list_input.append(signature[name].default      )

																for index, value in sorted_tuple_index_mapping:
																			SCREAMING_SNAKE_CASE     =       prepared_for_class[value]

																SCREAMING_SNAKE_CASE     =       tuple(a      )

																# Send to model
																SCREAMING_SNAKE_CASE     =       model(tuple_input[:-1]      )[0]

																self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1]      )







							@slow
							def  _UpperCAmelCase  (       self							:						str      )      ->       List[Any]:
										for model_name in TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE     =       TFMobileViTModel.from_pretrained(a      )
													self.assertIsNotNone(a      )





def        lowerCamelCase_     (     ):
			'''simple docstring'''
			SCREAMING_SNAKE_CASE     =       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""" )
			return image




@require_tf
class       UpperCAmelCase_   ( unittest.TestCase							):




							'''simple docstring'''



							@slow
							def  _UpperCAmelCase  (       self							:						Dict      )      ->       Dict:
										SCREAMING_SNAKE_CASE     =       TFMobileViTForImageClassification.from_pretrained("""apple/mobilevit-xx-small"""      )

										SCREAMING_SNAKE_CASE     =       MobileViTImageProcessor.from_pretrained("""apple/mobilevit-xx-small"""      )
										SCREAMING_SNAKE_CASE     =       prepare_img()
										SCREAMING_SNAKE_CASE     =       image_processor(images=a							,  return_tensors="""tf"""      )

										# forward pass
										SCREAMING_SNAKE_CASE     =       model(**a							,  training=a      )

										# verify the logits
										SCREAMING_SNAKE_CASE     =       tf.TensorShape((1, 1_000)      )
										self.assertEqual(outputs.logits.shape							,  a      )

										SCREAMING_SNAKE_CASE     =       tf.constant([-1.9364, -1.2327, -0.4653]      )

										tf.debugging.assert_near(outputs.logits[0, :3]							,  a							,  atol=1E-4							,  rtol=1E-0_4      )







							@slow
							def  _UpperCAmelCase  (       self							:						int      )      ->       Tuple:
										# `from_pt` will be removed
										SCREAMING_SNAKE_CASE     =       TFMobileViTForSemanticSegmentation.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""      )

										SCREAMING_SNAKE_CASE     =       MobileViTImageProcessor.from_pretrained("""apple/deeplabv3-mobilevit-xx-small"""      )

										SCREAMING_SNAKE_CASE     =       prepare_img()
										SCREAMING_SNAKE_CASE     =       image_processor(images=a							,  return_tensors="""tf"""      )

										# forward pass
										SCREAMING_SNAKE_CASE     =       model(inputs.pixel_values							,  training=a      )
										SCREAMING_SNAKE_CASE     =       outputs.logits

										# verify the logits
										SCREAMING_SNAKE_CASE     =       tf.TensorShape((1, 21, 32, 32)      )
										self.assertEqual(logits.shape							,  a      )

										SCREAMING_SNAKE_CASE     =       tf.constant(
										    [
										        [[6.9713, 6.9786, 7.2422], [7.2893, 7.2825, 7.4446], [7.6580, 7.8797, 7.9420]],
										        [[-10.6869, -10.3250, -10.3471], [-10.4228, -9.9868, -9.7132], [-11.0405, -11.0221, -10.7318]],
										        [[-3.3089, -2.8539, -2.6740], [-3.2706, -2.5621, -2.5108], [-3.2534, -2.6615, -2.6651]],
										    ]      )

										tf.debugging.assert_near(logits[0, :3, :3, :3]							,  a							,  rtol=1E-4							,  atol=1E-4      )






