




"""simple docstring"""

import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

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


if is_vision_available():
							from PIL import Image

							from transformers import ViTImageProcessor
class 	lowerCamelCase__			:


			def __init__(   self ,       SCREAMING_SNAKE_CASE ,       SCREAMING_SNAKE_CASE=13 ,       SCREAMING_SNAKE_CASE=30 ,       SCREAMING_SNAKE_CASE=2 ,       SCREAMING_SNAKE_CASE=3 ,       SCREAMING_SNAKE_CASE=True ,       SCREAMING_SNAKE_CASE=True ,       SCREAMING_SNAKE_CASE=32 ,       SCREAMING_SNAKE_CASE=5 ,       SCREAMING_SNAKE_CASE=4 ,       SCREAMING_SNAKE_CASE=37 ,       SCREAMING_SNAKE_CASE="gelu" ,       SCREAMING_SNAKE_CASE=0.1 ,       SCREAMING_SNAKE_CASE=0.1 ,       SCREAMING_SNAKE_CASE=10 ,       SCREAMING_SNAKE_CASE=0.02 ,       SCREAMING_SNAKE_CASE=None ,       SCREAMING_SNAKE_CASE=2 ,       ):



					"""simple docstring"""







					snake_case						:  List[str]    	= parent
					snake_case						:  int    	= batch_size
					snake_case						:  int    	= image_size
					snake_case						:  Any    	= patch_size
					snake_case						:  str    	= num_channels
					snake_case						:  List[Any]    	= is_training
					snake_case						:  Union[str, Any]    	= use_labels
					snake_case						:  Dict    	= hidden_size
					snake_case						:  List[str]    	= num_hidden_layers
					snake_case						:  Optional[Any]    	= num_attention_heads
					snake_case						:  List[str]    	= intermediate_size
					snake_case						:  Optional[int]    	= hidden_act
					snake_case						:  Dict    	= hidden_dropout_prob
					snake_case						:  List[Any]    	= attention_probs_dropout_prob
					snake_case						:  str    	= type_sequence_label_size
					snake_case						:  str    	= initializer_range
					snake_case						:  Dict    	= 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						:  List[str]    	= (image_size // patch_size) ** 2
					snake_case						:  Optional[Any]    	= num_patches + 1


			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					snake_case						:  List[str]    	= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

					snake_case						:  Union[str, Any]    	= 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 							lowerCamelCase_ (   self							):



					"""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=A__ ,       initializer_range=self.initializer_range ,       encoder_stride=self.encoder_stride ,       )


			def 							lowerCamelCase_ (   self ,       SCREAMING_SNAKE_CASE ,       SCREAMING_SNAKE_CASE ,       SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					snake_case						:  Optional[int]    	= ViTModel(config=A__							)
					model.to(A__							)
					model.eval()
					snake_case						:  int    	= model(A__							)
					self.parent.assertEqual(result.last_hidden_state.shape ,       (self.batch_size, self.seq_length, self.hidden_size)							)


			def 							lowerCamelCase_ (   self ,       SCREAMING_SNAKE_CASE ,       SCREAMING_SNAKE_CASE ,       SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					snake_case						:  Union[str, Any]    	= ViTForMaskedImageModeling(config=A__							)
					model.to(A__							)
					model.eval()
					snake_case						:  str    	= model(A__							)
					self.parent.assertEqual(
					    result.reconstruction.shape ,       (self.batch_size, self.num_channels, self.image_size, self.image_size)							)

					# test greyscale images
					snake_case						:  List[Any]    	= 1
					snake_case						:  str    	= ViTForMaskedImageModeling(A__							)
					model.to(A__							)
					model.eval()

					snake_case						:  Optional[Any]    	= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]							)
					snake_case						:  int    	= model(A__							)
					self.parent.assertEqual(result.reconstruction.shape ,       (self.batch_size, 1, self.image_size, self.image_size)							)


			def 							lowerCamelCase_ (   self ,       SCREAMING_SNAKE_CASE ,       SCREAMING_SNAKE_CASE ,       SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					snake_case						:  Any    	= self.type_sequence_label_size
					snake_case						:  List[str]    	= ViTForImageClassification(A__							)
					model.to(A__							)
					model.eval()
					snake_case						:  Tuple    	= model(A__ ,       labels=A__							)
					self.parent.assertEqual(result.logits.shape ,       (self.batch_size, self.type_sequence_label_size)							)

					# test greyscale images
					snake_case						:  int    	= 1
					snake_case						:  str    	= ViTForImageClassification(A__							)
					model.to(A__							)
					model.eval()

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



			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					snake_case						:  List[Any]    	= self.prepare_config_and_inputs()
					(
					    (
					    snake_case						
					)	, (
					    snake_case						
					)	, (
					    snake_case						
					)	, 						
					)						:  Optional[Any]    	= config_and_inputs
					snake_case						:  Union[str, Any]    	= {"pixel_values": pixel_values}
					return config, inputs_dict







@require_torch
class 	lowerCamelCase__			(     SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__ ,     unittest.TestCase	):
			a__       :					Tuple	    =   (
			    (
			        ViTModel,
			        ViTForImageClassification,
			        ViTForMaskedImageModeling,
			    )
			    if is_torch_available()
			    else ()
			)
			a__       :					List[Any]	    =   (
			    {'''feature-extraction''': ViTModel, '''image-classification''': ViTForImageClassification}
			    if is_torch_available()
			    else {}
			)
			a__       :					Dict	    =   True

			a__       :					Tuple	    =   False
			a__       :					Any	    =   False
			a__       :					str	    =   False


			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					snake_case						:  str    	= ViTModelTester(self							)
					snake_case						:  List[str]    	= ConfigTester(self ,       config_class=A__ ,       has_text_modality=A__ ,       hidden_size=37							)


			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					self.config_tester.run_common_tests()


			@unittest.skip(reason="ViT does not use inputs_embeds"							)
			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					pass


			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







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

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


			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







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

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

							snake_case						:  Optional[Any]    	= ["pixel_values"]
							self.assertListEqual(arg_names[:1] ,       A__							)


			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					snake_case						:  Dict    	= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*A__							)


			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					snake_case						:  Dict    	= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_image_modeling(*A__							)


			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					snake_case						:  List[str]    	= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*A__							)



			@slow
			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							snake_case						:  Tuple    	= ViTModel.from_pretrained(A__							)
							self.assertIsNotNone(A__							)







def      UpperCamelCase__   (						):
		snake_case						:  Optional[Any]    	= Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      )
		return image
@require_torch
@require_vision
class 	lowerCamelCase__			(     unittest.TestCase	):


			@cached_property
			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







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


			@slow
			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					snake_case						:  int    	= ViTForImageClassification.from_pretrained("google/vit-base-patch16-224"							).to(A__							)

					snake_case						:  Tuple    	= self.default_image_processor
					snake_case						:  Optional[int]    	= prepare_img()
					snake_case						:  str    	= image_processor(images=A__ ,       return_tensors="pt"							).to(A__							)

					# forward pass
					with torch.no_grad():
							snake_case						:  Optional[int]    	= model(**A__							)

					# verify the logits
					snake_case						:  str    	= torch.Size((1, 1_000)							)
					self.assertEqual(outputs.logits.shape ,       A__							)

					snake_case						:  Optional[Any]    	= torch.tensor([-0.27_44, 0.82_15, -0.08_36]							).to(A__							)

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


			@slow
			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







					snake_case						:  int    	= ViTModel.from_pretrained("facebook/dino-vits8"							).to(A__							)

					snake_case						:  Optional[Any]    	= ViTImageProcessor.from_pretrained("facebook/dino-vits8" ,       size=480							)
					snake_case						:  Optional[int]    	= prepare_img()
					snake_case						:  int    	= image_processor(images=A__ ,       return_tensors="pt"							)
					snake_case						:  List[Any]    	= inputs.pixel_values.to(A__							)

					# forward pass
					with torch.no_grad():
							snake_case						:  Dict    	= model(A__ ,       interpolate_pos_encoding=A__							)

					# verify the logits
					snake_case						:  Optional[int]    	= torch.Size((1, 3_601, 384)							)
					self.assertEqual(outputs.last_hidden_state.shape ,       A__							)

					snake_case						:  Dict    	= torch.tensor(
					    [[4.23_40, 4.39_06, -6.66_92], [4.54_63, 1.89_28, -6.72_57], [4.44_29, 0.84_96, -5.85_85]]							).to(A__							)

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



			@slow
			@require_accelerate
			@require_torch_gpu
			def 							lowerCamelCase_ (   self							):



					"""simple docstring"""







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

					snake_case						:  List[str]    	= prepare_img()
					snake_case						:  Union[str, Any]    	= image_processor(images=A__ ,       return_tensors="pt"							)
					snake_case						:  Union[str, Any]    	= inputs.pixel_values.to(A__							)

					# forward pass to make sure inference works in fp16
					with torch.no_grad():
							snake_case						:  List[Any]    	= model(A__							)



import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class 				lowercase		(					unittest.TestCase		):




							def  A__      (				self):
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()




							@property
							def  A__      (				self):
												lowercase							=						1
												lowercase							=						3
												lowercase							=						(3_2, 3_2)

												lowercase							=						floats_tensor((batch_size, num_channels) + sizes   ,rng=random.Random(0)).to(A__)
												return image




							@property
							def  A__      (				self):
												torch.manual_seed(0)
												lowercase							=						UNetaDConditionModel(
												    block_out_channels=(3_2, 6_4)   ,layers_per_block=2   ,sample_size=3_2   ,in_channels=4   ,out_channels=4   ,down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,cross_attention_dim=3_2   ,)
												return model




							@property
							def  A__      (				self):
												torch.manual_seed(0)
												lowercase							=						AutoencoderKL(
												    block_out_channels=[3_2, 6_4]   ,in_channels=3   ,out_channels=3   ,down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,latent_channels=4   ,)
												return model




							@property
							def  A__      (				self):
												torch.manual_seed(0)
												lowercase							=						RobertaSeriesConfig(
												    hidden_size=3_2   ,project_dim=3_2   ,intermediate_size=3_7   ,layer_norm_eps=1E-05   ,num_attention_heads=4   ,num_hidden_layers=5   ,pad_token_id=1   ,vocab_size=5_0_0_6   ,)
												return RobertaSeriesModelWithTransformation(A__)




							@property
							def  A__      (				self):
												def extract(*A__   ,**A__):







																	class 				lowercase		:




																								def __init__(				self):
																													lowercase							=						torch.ones([0])





																								def  A__      (				self   ,A__):
																													self.pixel_values.to(A__)
																													return self


																	return Out()

												return extract




							def  A__      (				self):
												lowercase							=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
												lowercase							=						self.dummy_cond_unet
												lowercase							=						PNDMScheduler(skip_prk_steps=A__)
												lowercase							=						self.dummy_vae
												lowercase							=						self.dummy_text_encoder
												lowercase							=						XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta''')
												lowercase							=						7_7

												lowercase							=						self.dummy_image.to(A__)
												lowercase							=						init_image / 2 + 0.5

												# make sure here that pndm scheduler skips prk
												lowercase							=						AltDiffusionImgaImgPipeline(
												    unet=A__   ,scheduler=A__   ,vae=A__   ,text_encoder=A__   ,tokenizer=A__   ,safety_checker=A__   ,feature_extractor=self.dummy_extractor   ,)
												lowercase							=						VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,do_normalize=A__)
												lowercase							=						alt_pipe.to(A__)
												alt_pipe.set_progress_bar_config(disable=A__)

												lowercase							=						'''A painting of a squirrel eating a burger'''
												lowercase							=						torch.Generator(device=A__).manual_seed(0)
												lowercase							=						alt_pipe(
												    [prompt]   ,generator=A__   ,guidance_scale=6.0   ,num_inference_steps=2   ,output_type='''np'''   ,image=A__   ,)

												lowercase							=						output.images

												lowercase							=						torch.Generator(device=A__).manual_seed(0)
												lowercase							=						alt_pipe(
												    [prompt]   ,generator=A__   ,guidance_scale=6.0   ,num_inference_steps=2   ,output_type='''np'''   ,image=A__   ,return_dict=A__   ,)[0]

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

												assert image.shape == (1, 3_2, 3_2, 3)
												lowercase							=						np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499])

												assert np.abs(image_slice.flatten() - expected_slice).max() < 5E-3
												assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 5E-3




							@unittest.skipIf(torch_device != '''cuda'''   ,'''This test requires a GPU''')
							def  A__      (				self):
												lowercase							=						self.dummy_cond_unet
												lowercase							=						PNDMScheduler(skip_prk_steps=A__)
												lowercase							=						self.dummy_vae
												lowercase							=						self.dummy_text_encoder
												lowercase							=						XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta''')
												lowercase							=						7_7

												lowercase							=						self.dummy_image.to(A__)

												# put models in fp16
												lowercase							=						unet.half()
												lowercase							=						vae.half()
												lowercase							=						bert.half()

												# make sure here that pndm scheduler skips prk
												lowercase							=						AltDiffusionImgaImgPipeline(
												    unet=A__   ,scheduler=A__   ,vae=A__   ,text_encoder=A__   ,tokenizer=A__   ,safety_checker=A__   ,feature_extractor=self.dummy_extractor   ,)
												lowercase							=						VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor   ,do_normalize=A__)
												lowercase							=						alt_pipe.to(A__)
												alt_pipe.set_progress_bar_config(disable=A__)

												lowercase							=						'''A painting of a squirrel eating a burger'''
												lowercase							=						torch.manual_seed(0)
												lowercase							=						alt_pipe(
												    [prompt]   ,generator=A__   ,num_inference_steps=2   ,output_type='''np'''   ,image=A__   ,).images

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





							@unittest.skipIf(torch_device != '''cuda'''   ,'''This test requires a GPU''')
							def  A__      (				self):
												lowercase							=						load_image(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/img2img/sketch-mountains-input.jpg''')
												# resize to resolution that is divisible by 8 but not 16 or 32
												lowercase							=						init_image.resize((7_6_0, 5_0_4))

												lowercase							=						'''BAAI/AltDiffusion'''
												lowercase							=						AltDiffusionImgaImgPipeline.from_pretrained(
												    A__   ,safety_checker=A__   ,)
												pipe.to(A__)
												pipe.set_progress_bar_config(disable=A__)
												pipe.enable_attention_slicing()

												lowercase							=						'''A fantasy landscape, trending on artstation'''

												lowercase							=						torch.manual_seed(0)
												lowercase							=						pipe(
												    prompt=A__   ,image=A__   ,strength=0.75   ,guidance_scale=7.5   ,generator=A__   ,output_type='''np'''   ,)
												lowercase							=						output.images[0]

												lowercase							=						image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

												assert image.shape == (5_0_4, 7_6_0, 3)
												lowercase							=						np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000])

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







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




							def  A__      (				self):
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()





							def  A__      (				self):
												lowercase							=						load_image(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
												    '''/img2img/sketch-mountains-input.jpg''')
												lowercase							=						init_image.resize((7_6_8, 5_1_2))
												lowercase							=						load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy''')

												lowercase							=						'''BAAI/AltDiffusion'''
												lowercase							=						AltDiffusionImgaImgPipeline.from_pretrained(
												    A__   ,safety_checker=A__   ,)
												pipe.to(A__)
												pipe.set_progress_bar_config(disable=A__)
												pipe.enable_attention_slicing()

												lowercase							=						'''A fantasy landscape, trending on artstation'''

												lowercase							=						torch.manual_seed(0)
												lowercase							=						pipe(
												    prompt=A__   ,image=A__   ,strength=0.75   ,guidance_scale=7.5   ,generator=A__   ,output_type='''np'''   ,)
												lowercase							=						output.images[0]

												assert image.shape == (5_1_2, 7_6_8, 3)
												# img2img is flaky across GPUs even in fp32, so using MAE here
												assert np.abs(expected_image - image).max() < 1E-2




