





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 SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






						def UpperCAmelCase_				(      self      :      int	)							->    Tuple:







							"""simple docstring"""


							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






						@property
						def UpperCAmelCase_				(      self      :      Optional[int]	)							->    Tuple:







							"""simple docstring"""


							snake_case_		:			Tuple											=	1
							snake_case_		:			Optional[Any]											=	3
							snake_case_		:			Optional[int]											=	(32, 32)

							snake_case_		:			int											=	floats_tensor((batch_size, num_channels) + sizes		,   rng=random.Random(0	)	).to(_A	)
							return image






						@property
						def UpperCAmelCase_				(      self      :      Optional[int]	)							->    Union[str, Any]:







							"""simple docstring"""


							torch.manual_seed(0	)
							snake_case_		:			Dict											=	UNetaDConditionModel(
							    block_out_channels=(32, 64)		,   layers_per_block=2		,   sample_size=32		,   in_channels=4		,   out_channels=4		,   down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')		,   up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')		,   cross_attention_dim=32		,   )
							return model






						@property
						def UpperCAmelCase_				(      self      :      Union[str, Any]	)							->    Optional[Any]:







							"""simple docstring"""


							torch.manual_seed(0	)
							snake_case_		:			Dict											=	AutoencoderKL(
							    block_out_channels=[32, 64]		,   in_channels=3		,   out_channels=3		,   down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']		,   up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']		,   latent_channels=4		,   )
							return model






						@property
						def UpperCAmelCase_				(      self      :      Optional[Any]	)							->    List[Any]:







							"""simple docstring"""


							torch.manual_seed(0	)
							snake_case_		:			Any											=	RobertaSeriesConfig(
							    hidden_size=32		,   project_dim=32		,   intermediate_size=37		,   layer_norm_eps=1E-05		,   num_attention_heads=4		,   num_hidden_layers=5		,   pad_token_id=1		,   vocab_size=5006		,   )
							return RobertaSeriesModelWithTransformation(_A	)






						@property
						def UpperCAmelCase_				(      self      :      int	)							->    Dict:







							"""simple docstring"""


							def extract(*_A      :      Any		,   **_A      :      Tuple	):
								class SCREAMING_SNAKE_CASE_						:






														def __init__(      self      :      Dict	)							->    int:







															"""simple docstring"""


															snake_case_		:			List[str]											=	torch.ones([0]	)






														def UpperCAmelCase_				(      self      :      Any		,   _A      :      Optional[Any]	)							->    Optional[int]:







															"""simple docstring"""


															self.pixel_values.to(_A	)
															return self
								return Out()

							return extract






						def UpperCAmelCase_				(      self      :      List[str]	)							->    Dict:







							"""simple docstring"""


							snake_case_		:			Optional[int]											=	'cpu'  # ensure determinism for the device-dependent torch.Generator
							snake_case_		:			int											=	self.dummy_cond_unet
							snake_case_		:			Union[str, Any]											=	PNDMScheduler(skip_prk_steps=_A	)
							snake_case_		:			List[Any]											=	self.dummy_vae
							snake_case_		:			Tuple											=	self.dummy_text_encoder
							snake_case_		:			Optional[int]											=	XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'	)
							snake_case_		:			Union[str, Any]											=	77

							snake_case_		:			Tuple											=	self.dummy_image.to(_A	)
							snake_case_		:			Optional[Any]											=	init_image / 2 + 0.5

							# make sure here that pndm scheduler skips prk
							snake_case_		:			Optional[int]											=	AltDiffusionImgaImgPipeline(
							    unet=_A		,   scheduler=_A		,   vae=_A		,   text_encoder=_A		,   tokenizer=_A		,   safety_checker=_A		,   feature_extractor=self.dummy_extractor		,   )
							snake_case_		:			int											=	VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		,   do_normalize=_A	)
							snake_case_		:			Dict											=	alt_pipe.to(_A	)
							alt_pipe.set_progress_bar_config(disable=_A	)

							snake_case_		:			Union[str, Any]											=	'A painting of a squirrel eating a burger'
							snake_case_		:			str											=	torch.Generator(device=_A	).manual_seed(0	)
							snake_case_		:			Dict											=	alt_pipe(
							    [prompt]		,   generator=_A		,   guidance_scale=6.0		,   num_inference_steps=2		,   output_type='np'		,   image=_A		,   )

							snake_case_		:			str											=	output.images

							snake_case_		:			Optional[Any]											=	torch.Generator(device=_A	).manual_seed(0	)
							snake_case_		:			Optional[int]											=	alt_pipe(
							    [prompt]		,   generator=_A		,   guidance_scale=6.0		,   num_inference_steps=2		,   output_type='np'		,   image=_A		,   return_dict=_A		,   )[0]

							snake_case_		:			List[str]											=	image[0, -3:, -3:, -1]
							snake_case_		:			List[str]											=	image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 32, 32, 3)
							snake_case_		:			List[Any]											=	np.array([0.4_4_2_7, 0.3_7_3_1, 0.4_2_4_9, 0.4_9_4_1, 0.4_5_4_6, 0.4_1_4_8, 0.4_1_9_3, 0.4_6_6_6, 0.4_4_9_9]	)

							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 UpperCAmelCase_				(      self      :      Dict	)							->    Tuple:







							"""simple docstring"""


							snake_case_		:			Optional[Any]											=	self.dummy_cond_unet
							snake_case_		:			List[Any]											=	PNDMScheduler(skip_prk_steps=_A	)
							snake_case_		:			Optional[int]											=	self.dummy_vae
							snake_case_		:			Dict											=	self.dummy_text_encoder
							snake_case_		:			int											=	XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'	)
							snake_case_		:			Tuple											=	77

							snake_case_		:			Tuple											=	self.dummy_image.to(_A	)

							# put models in fp16
							snake_case_		:			Dict											=	unet.half()
							snake_case_		:			Any											=	vae.half()
							snake_case_		:			Any											=	bert.half()

							# make sure here that pndm scheduler skips prk
							snake_case_		:			int											=	AltDiffusionImgaImgPipeline(
							    unet=_A		,   scheduler=_A		,   vae=_A		,   text_encoder=_A		,   tokenizer=_A		,   safety_checker=_A		,   feature_extractor=self.dummy_extractor		,   )
							snake_case_		:			Dict											=	VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		,   do_normalize=_A	)
							snake_case_		:			List[str]											=	alt_pipe.to(_A	)
							alt_pipe.set_progress_bar_config(disable=_A	)

							snake_case_		:			Dict											=	'A painting of a squirrel eating a burger'
							snake_case_		:			Optional[Any]											=	torch.manual_seed(0	)
							snake_case_		:			Union[str, Any]											=	alt_pipe(
							    [prompt]		,   generator=_A		,   num_inference_steps=2		,   output_type='np'		,   image=_A		,   ).images

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






						@unittest.skipIf(torch_device != 'cuda'		,   'This test requires a GPU'	)
						def UpperCAmelCase_				(      self      :      Optional[int]	)							->    List[str]:







							"""simple docstring"""


							snake_case_		:			List[Any]											=	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
							snake_case_		:			List[str]											=	init_image.resize((760, 504)	)

							snake_case_		:			Optional[Any]											=	'BAAI/AltDiffusion'
							snake_case_		:			List[str]											=	AltDiffusionImgaImgPipeline.from_pretrained(
							    _A		,   safety_checker=_A		,   )
							pipe.to(_A	)
							pipe.set_progress_bar_config(disable=_A	)
							pipe.enable_attention_slicing()

							snake_case_		:			List[Any]											=	'A fantasy landscape, trending on artstation'

							snake_case_		:			Dict											=	torch.manual_seed(0	)
							snake_case_		:			Dict											=	pipe(
							    prompt=_A		,   image=_A		,   strength=0.7_5		,   guidance_scale=7.5		,   generator=_A		,   output_type='np'		,   )
							snake_case_		:			Dict											=	output.images[0]

							snake_case_		:			List[Any]											=	image[255:258, 383:386, -1]

							assert image.shape == (504, 760, 3)
							snake_case_		:			int											=	np.array([0.9_3_5_8, 0.9_3_9_7, 0.9_5_9_9, 0.9_9_0_1, 1.0_0_0_0, 1.0_0_0_0, 0.9_8_8_2, 1.0_0_0_0, 1.0_0_0_0]	)

							assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2
@slow
@require_torch_gpu
class SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






						def UpperCAmelCase_				(      self      :      Tuple	)							->    Any:







							"""simple docstring"""


							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






						def UpperCAmelCase_				(      self      :      Tuple	)							->    Dict:







							"""simple docstring"""


							snake_case_		:			List[Any]											=	load_image(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/img2img/sketch-mountains-input.jpg'	)
							snake_case_		:			str											=	init_image.resize((768, 512)	)
							snake_case_		:			Optional[int]											=	load_numpy(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy'	)

							snake_case_		:			Union[str, Any]											=	'BAAI/AltDiffusion'
							snake_case_		:			Tuple											=	AltDiffusionImgaImgPipeline.from_pretrained(
							    _A		,   safety_checker=_A		,   )
							pipe.to(_A	)
							pipe.set_progress_bar_config(disable=_A	)
							pipe.enable_attention_slicing()

							snake_case_		:			Tuple											=	'A fantasy landscape, trending on artstation'

							snake_case_		:			str											=	torch.manual_seed(0	)
							snake_case_		:			Union[str, Any]											=	pipe(
							    prompt=_A		,   image=_A		,   strength=0.7_5		,   guidance_scale=7.5		,   generator=_A		,   output_type='np'		,   )
							snake_case_		:			Optional[int]											=	output.images[0]

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








import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
						import jax

						from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel
class SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






						def __init__(      self      :      Dict		,   _A      :      List[str]		,   _A      :      Dict=13		,   _A      :      str=30		,   _A      :      Any=2		,   _A      :      Dict=3		,   _A      :      Optional[Any]=True		,   _A      :      Tuple=True		,   _A      :      List[str]=32		,   _A      :      int=5		,   _A      :      Optional[int]=4		,   _A      :      Dict=37		,   _A      :      Tuple="gelu"		,   _A      :      Dict=0.1		,   _A      :      str=0.1		,   _A      :      int=10		,   _A      :      Union[str, Any]=0.0_2		,   )							->    int:







							"""simple docstring"""


							snake_case_		:			int											=	parent
							snake_case_		:			str											=	batch_size
							snake_case_		:			List[str]											=	image_size
							snake_case_		:			Tuple											=	patch_size
							snake_case_		:			str											=	num_channels
							snake_case_		:			List[str]											=	is_training
							snake_case_		:			List[str]											=	use_labels
							snake_case_		:			str											=	hidden_size
							snake_case_		:			Union[str, Any]											=	num_hidden_layers
							snake_case_		:			Any											=	num_attention_heads
							snake_case_		:			List[Any]											=	intermediate_size
							snake_case_		:			Any											=	hidden_act
							snake_case_		:			List[Any]											=	hidden_dropout_prob
							snake_case_		:			List[Any]											=	attention_probs_dropout_prob
							snake_case_		:			Tuple											=	type_sequence_label_size
							snake_case_		:			Any											=	initializer_range

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






						def UpperCAmelCase_				(      self      :      Dict	)							->    List[str]:







							"""simple docstring"""


							snake_case_		:			int											=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

							snake_case_		:			Optional[Any]											=	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		,   )

							return config, pixel_values






						def UpperCAmelCase_				(      self      :      int		,   _A      :      Union[str, Any]		,   _A      :      List[Any]	)							->    Dict:







							"""simple docstring"""


							snake_case_		:			Union[str, Any]											=	FlaxViTModel(config=_A	)
							snake_case_		:			Optional[Any]											=	model(_A	)
							# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
							snake_case_		:			Union[str, Any]											=	(self.image_size, self.image_size)
							snake_case_		:			int											=	(self.patch_size, self.patch_size)
							snake_case_		:			Optional[Any]											=	(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
							self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, num_patches + 1, self.hidden_size)	)






						def UpperCAmelCase_				(      self      :      str		,   _A      :      Any		,   _A      :      List[str]	)							->    Dict:







							"""simple docstring"""


							snake_case_		:			Dict											=	self.type_sequence_label_size
							snake_case_		:			Any											=	FlaxViTForImageClassification(config=_A	)
							snake_case_		:			Optional[int]											=	model(_A	)
							self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.type_sequence_label_size)	)

							# test greyscale images
							snake_case_		:			List[str]											=	1
							snake_case_		:			Dict											=	FlaxViTForImageClassification(_A	)

							snake_case_		:			Optional[int]											=	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
							snake_case_		:			Any											=	model(_A	)






						def UpperCAmelCase_				(      self      :      Any	)							->    Any:







							"""simple docstring"""


							snake_case_		:			Tuple											=	self.prepare_config_and_inputs()
							(
							    (
							    snake_case_
							)						,(
							    snake_case_
							)						,
							)		:			Any											=	config_and_inputs
							snake_case_		:			Tuple											=	{'pixel_values': pixel_values}
							return config, inputs_dict
@require_flax
class SCREAMING_SNAKE_CASE_						(					snake_case_       ,				unittest.TestCase					):
						__magic_name__:			Tuple       =       (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()






						def UpperCAmelCase_				(      self      :      List[str]	)							->    None:







							"""simple docstring"""


							snake_case_		:			Union[str, Any]											=	FlaxViTModelTester(self	)
							snake_case_		:			Any											=	ConfigTester(self		,   config_class=_A		,   has_text_modality=_A		,   hidden_size=37	)






						def UpperCAmelCase_				(      self      :      Any	)							->    Dict:







							"""simple docstring"""


							self.config_tester.run_common_tests()






						def UpperCAmelCase_				(      self      :      Any	)							->    Optional[Any]:







							"""simple docstring"""


							snake_case_		:			str											=	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_A	)






						def UpperCAmelCase_				(      self      :      List[Any]	)							->    Union[str, Any]:







							"""simple docstring"""


							snake_case_		:			Tuple											=	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_A	)






						def UpperCAmelCase_				(      self      :      int	)							->    List[str]:







							"""simple docstring"""


							snake_case_						,snake_case_		:			Tuple											=	self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
								snake_case_		:			str											=	model_class(_A	)
								snake_case_		:			Dict											=	inspect.signature(model.__call__	)
								# signature.parameters is an OrderedDict => so arg_names order is deterministic
								snake_case_		:			Optional[Any]											=	[*signature.parameters.keys()]

								snake_case_		:			Dict											=	['pixel_values']
								self.assertListEqual(arg_names[:1]		,   _A	)






						def UpperCAmelCase_				(      self      :      Optional[Any]	)							->    str:







							"""simple docstring"""


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

							for model_class in self.all_model_classes:
								with self.subTest(model_class.__name__	):
									snake_case_		:			Dict											=	self._prepare_for_class(_A		,   _A	)
									snake_case_		:			Tuple											=	model_class(_A	)

									@jax.jit
									def model_jitted(_A      :      Optional[int]		,   **_A      :      Any	):
										return model(pixel_values=_A		,   **_A	)

									with self.subTest('JIT Enabled'	):
										snake_case_		:			Optional[Any]											=	model_jitted(**_A	).to_tuple()

									with self.subTest('JIT Disabled'	):
										with jax.disable_jit():
											snake_case_		:			Optional[Any]											=	model_jitted(**_A	).to_tuple()

									self.assertEqual(len(_A	)		,   len(_A	)	)
									for jitted_output, output in zip(_A		,   _A	):
										self.assertEqual(jitted_output.shape		,   output.shape	)






						@slow
						def UpperCAmelCase_				(      self      :      List[str]	)							->    Any:







							"""simple docstring"""


							for model_class_name in self.all_model_classes:
								snake_case_		:			List[Any]											=	model_class_name.from_pretrained('google/vit-base-patch16-224'	)
								snake_case_		:			Optional[Any]											=	model(np.ones((1, 3, 224, 224)	)	)
								self.assertIsNotNone(_A	)

