



import gc
import unittest

from diffusers import FlaxStableDiffusionInpaintPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
							import jax
							import jax.numpy as jnp
							from flax.jax_utils import replicate
							from flax.training.common_utils import shard


@slow
@require_flax
class   __lowerCAmelCase  (		unittest.TestCase   ):
					"""simple docstring"""






					def      _a				(  self       :							Union[str, Any]       ):






								"""simple docstring"""


								super().tearDown()
								gc.collect()



					def      _a				(  self       :							str       ):






								"""simple docstring"""


								A__	  =     load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/sd2-inpaint/init_image.png'       )
								A__	  =     load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png'       )

								A__	  =     'xvjiarui/stable-diffusion-2-inpainting'
								A__  ,			A__	  =     FlaxStableDiffusionInpaintPipeline.from_pretrained(_snake_case					,       safety_checker=_snake_case       )

								A__	  =     'Face of a yellow cat, high resolution, sitting on a park bench'

								A__	  =     jax.random.PRNGKey(0       )
								A__	  =     50

								A__	  =     jax.device_count()
								A__	  =     num_samples * [prompt]
								A__	  =     num_samples * [init_image]
								A__	  =     num_samples * [mask_image]
								A__  ,			A__  ,			A__	  =     pipeline.prepare_inputs(_snake_case					,       _snake_case					,       _snake_case       )

								# shard inputs and rng
								A__	  =     replicate(_snake_case       )
								A__	  =     jax.random.split(_snake_case					,       jax.device_count()       )
								A__	  =     shard(_snake_case       )
								A__	  =     shard(_snake_case       )
								A__	  =     shard(_snake_case       )

								A__	  =     pipeline(
								    _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       _snake_case					,       jit=_snake_case       )

								A__	  =     output.images.reshape(_snake_case					,       5_12					,       5_12					,       3       )

								A__	  =     images[0, 2_53:2_56, 2_53:2_56, -1]

								A__	  =     jnp.asarray(jax.device_get(image_slice.flatten()       )       )
								A__	  =     jnp.array(
								    [0.361_1307, 0.3764_9736, 0.375_7408, 0.3821_3953, 0.3929_5167, 0.384_1631, 0.4155_4978, 0.413_7475, 0.421_7084]       )
								print(F'''output_slice: {output_slice}'''       )
								assert jnp.abs(output_slice - expected_slice       ).max() < 1E-2








'''simple docstring'''
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow


if is_torch_available():
		import torch

		from transformers import XLMRobertaModel
@require_sentencepiece
@require_tokenizers
@require_torch
class   lowercase      (				unittest.TestCase   ):

		"""simple docstring"""

		@slow
		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							UpperCamelCase__  :Optional[Any]				   =					XLMRobertaModel.from_pretrained('''xlm-roberta-base'''				)
							UpperCamelCase__  :List[str]				   =					torch.tensor([[0, 581, 10269, 83, 99942, 136, 60742, 23, 70, 80583, 18276, 2]]				)
							# The dog is cute and lives in the garden house

							UpperCamelCase__  :Union[str, Any]				   =					torch.Size((1, 12, 768)				)  # batch_size, sequence_length, embedding_vector_dim
							UpperCamelCase__  :Union[str, Any]				   =					torch.tensor(
							    [[-0.0101, 0.1218, -0.0803, 0.0801, 0.1327, 0.0776, -0.1215, 0.2383, 0.3338, 0.3106, 0.0300, 0.0252]]				)
							#  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.base')
							#  xlmr.eval()
							#  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
							with torch.no_grad():
												UpperCamelCase__  :Any				   =					model(UpperCamelCase_				)['''last_hidden_state'''].detach()
							self.assertEqual(output.shape       ,  UpperCamelCase_				)
							# compare the actual values for a slice of last dim
							self.assertTrue(torch.allclose(output[:, :, -1]       ,  UpperCamelCase_       ,  atol=1e-3				)				)




		@slow
		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							UpperCamelCase__  :int				   =					XLMRobertaModel.from_pretrained('''xlm-roberta-large'''				)
							UpperCamelCase__  :List[str]				   =					torch.tensor([[0, 581, 10269, 83, 99942, 136, 60742, 23, 70, 80583, 18276, 2]]				)
							# The dog is cute and lives in the garden house

							UpperCamelCase__  :Any				   =					torch.Size((1, 12, 1024)				)  # batch_size, sequence_length, embedding_vector_dim
							UpperCamelCase__  :List[Any]				   =					torch.tensor(
							    [[-0.0699, -0.0318, 0.0705, -0.1241, 0.0999, -0.0520, 0.1004, -0.1838, -0.4704, 0.1437, 0.0821, 0.0126]]				)
							#  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.large')
							#  xlmr.eval()
							#  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
							with torch.no_grad():
												UpperCamelCase__  :Optional[int]				   =					model(UpperCamelCase_				)['''last_hidden_state'''].detach()
							self.assertEqual(output.shape       ,  UpperCamelCase_				)
							# compare the actual values for a slice of last dim
							self.assertTrue(torch.allclose(output[:, :, -1]       ,  UpperCamelCase_       ,  atol=1e-3				)				)