






"""simple docstring"""






import gc
import unittest

from parameterized import parameterized

from diffusers import FlaxUNetaDConditionModel
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import load_hf_numpy, require_flax, slow


if is_flax_available():
					import jax
					import jax.numpy as jnp





@slow
@require_flax
class       snake_case_       (    unittest.TestCase       ):



	"""simple docstring"""

	def       _UpperCAmelCase						(      self  ,					__a  ,					__a						):


					"""simple docstring"""

					return f'''gaussian_noise_s={seed}_shape={'_'.join([str(__a						) for s in shape]						)}.npy'''


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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


	def       _UpperCAmelCase						(      self  ,					__a=0  ,					__a=(4, 4, 64, 64)  ,					__a=False						):


					"""simple docstring"""

					A__							 =       jnp.bfloataa if fpaa else jnp.floataa
					A__							 =       jnp.array(load_hf_numpy(self.get_file_format(__a  ,					__a						)						)  ,					dtype=__a						)
					return image


	def       _UpperCAmelCase						(      self  ,					__a=False  ,					__a="CompVis/stable-diffusion-v1-4"						):


					"""simple docstring"""

					A__							 =       jnp.bfloataa if fpaa else jnp.floataa
					A__							 =       'bf16' if fpaa else None

					A__ ,			A__							 =       FlaxUNetaDConditionModel.from_pretrained(
					    __a  ,					subfolder='unet'  ,					dtype=__a  ,					revision=__a						)
					return model, params


	def       _UpperCAmelCase						(      self  ,					__a=0  ,					__a=(4, 77, 768)  ,					__a=False						):


					"""simple docstring"""

					A__							 =       jnp.bfloataa if fpaa else jnp.floataa
					A__							 =       jnp.array(load_hf_numpy(self.get_file_format(__a  ,					__a						)						)  ,					dtype=__a						)
					return hidden_states


	@parameterized.expand(
	    [
	        # fmt: off
	        [83, 4, [-0.2323, -0.1304, 0.0813, -0.3093, -0.0919, -0.1571, -0.1125, -0.5806]],
	        [17, 0.55, [-0.0831, -0.2443, 0.0901, -0.0919, 0.3396, 0.0103, -0.3743, 0.0701]],
	        [8, 0.89, [-0.4863, 0.0859, 0.0875, -0.1658, 0.9199, -0.0114, 0.4839, 0.4639]],
	        [3, 1000, [-0.5649, 0.2402, -0.5518, 0.1248, 1.1328, -0.2443, -0.0325, -1.0078]],
	        # fmt: on
	    ]						)
	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a						):


					"""simple docstring"""

					A__ ,			A__							 =       self.get_unet_model(model_id='CompVis/stable-diffusion-v1-4'  ,					fpaa=__a						)
					A__							 =       self.get_latents(__a  ,					fpaa=__a						)
					A__							 =       self.get_encoder_hidden_states(__a  ,					fpaa=__a						)

					A__							 =       model.apply(
					    {'params': params}  ,					__a  ,					jnp.array(__a  ,					dtype=jnp.intaa						)  ,					encoder_hidden_states=__a  ,					).sample

					assert sample.shape == latents.shape

					A__							 =       jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())						)  ,					dtype=jnp.floataa						)
					A__							 =       jnp.array(__a  ,					dtype=jnp.floataa						)

					# Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, in the same hardware
					assert jnp.allclose(__a  ,					__a  ,					atol=1E-2						)






	@parameterized.expand(
	    [
	        # fmt: off
	        [83, 4, [0.1514, 0.0807, 0.1624, 0.1016, -0.1896, 0.0263, 0.0677, 0.2310]],
	        [17, 0.55, [0.1164, -0.0216, 0.0170, 0.1589, -0.3120, 0.1005, -0.0581, -0.1458]],
	        [8, 0.89, [-0.1758, -0.0169, 0.1004, -0.1411, 0.1312, 0.1103, -0.1996, 0.2139]],
	        [3, 1000, [0.1214, 0.0352, -0.0731, -0.1562, -0.0994, -0.0906, -0.2340, -0.0539]],
	        # fmt: on
	    ]						)
	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a						):


					"""simple docstring"""

					A__ ,			A__							 =       self.get_unet_model(model_id='stabilityai/stable-diffusion-2'  ,					fpaa=__a						)
					A__							 =       self.get_latents(__a  ,					shape=(4, 4, 96, 96)  ,					fpaa=__a						)
					A__							 =       self.get_encoder_hidden_states(__a  ,					shape=(4, 77, 1024)  ,					fpaa=__a						)

					A__							 =       model.apply(
					    {'params': params}  ,					__a  ,					jnp.array(__a  ,					dtype=jnp.intaa						)  ,					encoder_hidden_states=__a  ,					).sample

					assert sample.shape == latents.shape

					A__							 =       jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())						)  ,					dtype=jnp.floataa						)
					A__							 =       jnp.array(__a  ,					dtype=jnp.floataa						)

					# Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, on the same hardware
					assert jnp.allclose(__a  ,					__a  ,					atol=1E-2						)










"""simple docstring"""






import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
					import torch

if is_vision_available():
					from PIL import Image

					from transformers import PoolFormerImageProcessor





class       snake_case_       (    unittest.TestCase       ):



	"""simple docstring"""

	def __init__(      self  ,					__a  ,					__a=7  ,					__a=3  ,					__a=30  ,					__a=400  ,					__a=True  ,					__a=None  ,					__a=0.9  ,					__a=None  ,					__a=True  ,					__a=[0.5, 0.5, 0.5]  ,					__a=[0.5, 0.5, 0.5]  ,					):


					"""simple docstring"""

					A__							 =       size if size is not None else {'shortest_edge': 30}
					A__							 =       crop_size if crop_size is not None else {'height': 30, 'width': 30}
					A__							 =       parent
					A__							 =       batch_size
					A__							 =       num_channels
					A__							 =       min_resolution
					A__							 =       max_resolution
					A__							 =       do_resize_and_center_crop
					A__							 =       size
					A__							 =       crop_pct
					A__							 =       crop_size
					A__							 =       do_normalize
					A__							 =       image_mean
					A__							 =       image_std






	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					return {
					    "size": self.size,
					    "do_resize_and_center_crop": self.do_resize_and_center_crop,
					    "crop_pct": self.crop_pct,
					    "crop_size": self.crop_size,
					    "do_normalize": self.do_normalize,
					    "image_mean": self.image_mean,
					    "image_std": self.image_std,
					}





@require_torch
@require_vision
class       snake_case_       (    _lowerCamelCase    ,							unittest.TestCase       ):



	"""simple docstring"""

	SCREAMING_SNAKE_CASE_:							List[Any]   			=					PoolFormerImageProcessor if is_vision_available() else None


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       PoolFormerImageProcessingTester(self						)


	@property
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					return self.image_processor_tester.prepare_image_processor_dict()


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class(**self.image_processor_dict						)
					self.assertTrue(hasattr(__a  ,					'do_resize_and_center_crop'						)						)
					self.assertTrue(hasattr(__a  ,					'size'						)						)
					self.assertTrue(hasattr(__a  ,					'crop_pct'						)						)
					self.assertTrue(hasattr(__a  ,					'do_normalize'						)						)
					self.assertTrue(hasattr(__a  ,					'image_mean'						)						)
					self.assertTrue(hasattr(__a  ,					'image_std'						)						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class.from_dict(self.image_processor_dict						)
					self.assertEqual(image_processor.size  ,					{'shortest_edge': 30}						)
					self.assertEqual(image_processor.crop_size  ,					{'height': 30, 'width': 30}						)

					A__							 =       self.image_processing_class.from_dict(self.image_processor_dict  ,					size=42  ,					crop_size=84						)
					self.assertEqual(image_processor.size  ,					{'shortest_edge': 42}						)
					self.assertEqual(image_processor.crop_size  ,					{'height': 84, 'width': 84}						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					pass


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class(**self.image_processor_dict						)
					# create random PIL images
					A__							 =       prepare_image_inputs(self.image_processor_tester  ,					equal_resolution=__a						)
					for image in image_inputs:
									self.assertIsInstance(__a  ,					Image.Image						)

					# Test not batched input
					A__							 =       image_processing(image_inputs[0]  ,					return_tensors='pt'						).pixel_values

					self.assertEqual(
					    encoded_images.shape  ,					(
					        1,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)

					# Test batched
					A__							 =       image_processing(__a  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        self.image_processor_tester.batch_size,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class(**self.image_processor_dict						)
					# create random numpy tensors
					A__							 =       prepare_image_inputs(self.image_processor_tester  ,					equal_resolution=__a  ,					numpify=__a						)
					for image in image_inputs:
									self.assertIsInstance(__a  ,					np.ndarray						)

					# Test not batched input
					A__							 =       image_processing(image_inputs[0]  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        1,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)

					# Test batched
					A__							 =       image_processing(__a  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        self.image_processor_tester.batch_size,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)






	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.image_processing_class(**self.image_processor_dict						)
					# create random PyTorch tensors
					A__							 =       prepare_image_inputs(self.image_processor_tester  ,					equal_resolution=__a  ,					torchify=__a						)
					for image in image_inputs:
									self.assertIsInstance(__a  ,					torch.Tensor						)

					# Test not batched input
					A__							 =       image_processing(image_inputs[0]  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        1,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)

					# Test batched
					A__							 =       image_processing(__a  ,					return_tensors='pt'						).pixel_values
					self.assertEqual(
					    encoded_images.shape  ,					(
					        self.image_processor_tester.batch_size,
					        self.image_processor_tester.num_channels,
					        self.image_processor_tester.crop_size['height'],
					        self.image_processor_tester.crop_size['width'],
					    )  ,					)


