





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






			def 			lowerCamelCase  (self					,							__magic_name__					,							__magic_name__    )     ->   Optional[int]:


							'''simple docstring'''



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






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


							'''simple docstring'''



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






			def 			lowerCamelCase  (self					,							__magic_name__=0					,							__magic_name__=(4, 4, 64, 64)					,							__magic_name__=False    )     ->   str:


							'''simple docstring'''



							snake_case_						:    List[str]				=  jnp.bfloataa if fpaa else jnp.floataa
							snake_case_						:    Tuple				=  jnp.array(load_hf_numpy(self.get_file_format(SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_    )    )					,							dtype=SCREAMING_SNAKE_CASE_    )
							return image






			def 			lowerCamelCase  (self					,							__magic_name__=False					,							__magic_name__="CompVis/stable-diffusion-v1-4"    )     ->   Dict:


							'''simple docstring'''



							snake_case_						:    Union[str, Any]				=  jnp.bfloataa if fpaa else jnp.floataa
							snake_case_						:    Optional[Any]				=  'bf16' if fpaa else None

							snake_case_						:    str				=  FlaxUNetaDConditionModel.from_pretrained(
							    SCREAMING_SNAKE_CASE_					,							subfolder='''unet'''					,							dtype=SCREAMING_SNAKE_CASE_					,							revision=SCREAMING_SNAKE_CASE_    )
							return model, params






			def 			lowerCamelCase  (self					,							__magic_name__=0					,							__magic_name__=(4, 77, 768)					,							__magic_name__=False    )     ->   Union[str, Any]:


							'''simple docstring'''



							snake_case_						:    Any				=  jnp.bfloataa if fpaa else jnp.floataa
							snake_case_						:    Optional[Any]				=  jnp.array(load_hf_numpy(self.get_file_format(SCREAMING_SNAKE_CASE_					,							SCREAMING_SNAKE_CASE_    )    )					,							dtype=SCREAMING_SNAKE_CASE_    )
							return hidden_states






			@parameterized.expand(
			    [
			        # fmt: off
			        [83, 4, [-0.2_323, -0.1_304, 0.0_813, -0.3_093, -0.0_919, -0.1_571, -0.1_125, -0.5_806]],
			        [17, 0.55, [-0.0_831, -0.2_443, 0.0_901, -0.0_919, 0.3_396, 0.0_103, -0.3_743, 0.0_701]],
			        [8, 0.89, [-0.4_863, 0.0_859, 0.0_875, -0.1_658, 0.9_199, -0.0_114, 0.4_839, 0.4_639]],
			        [3, 1000, [-0.5_649, 0.2_402, -0.5_518, 0.1_248, 1.1_328, -0.2_443, -0.0_325, -1.0_078]],
			        # fmt: on
			    ]    )
			def 			lowerCamelCase  (self					,							__magic_name__					,							__magic_name__					,							__magic_name__    )     ->   Any:


							'''simple docstring'''



							snake_case_						:    Tuple				=  self.get_unet_model(model_id='''CompVis/stable-diffusion-v1-4'''					,							fpaa=SCREAMING_SNAKE_CASE_    )
							snake_case_						:    Any				=  self.get_latents(SCREAMING_SNAKE_CASE_					,							fpaa=SCREAMING_SNAKE_CASE_    )
							snake_case_						:    int				=  self.get_encoder_hidden_states(SCREAMING_SNAKE_CASE_					,							fpaa=SCREAMING_SNAKE_CASE_    )

							snake_case_						:    List[Any]				=  model.apply(
							    {'''params''': params}					,							SCREAMING_SNAKE_CASE_					,							jnp.array(SCREAMING_SNAKE_CASE_					,							dtype=jnp.intaa    )					,							encoder_hidden_states=SCREAMING_SNAKE_CASE_					,							).sample

							assert sample.shape == latents.shape

							snake_case_						:    Any				=  jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())    )					,							dtype=jnp.floataa    )
							snake_case_						:    Optional[Any]				=  jnp.array(SCREAMING_SNAKE_CASE_					,							dtype=jnp.floataa    )

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






			@parameterized.expand(
			    [
			        # fmt: off
			        [83, 4, [0.1_514, 0.0_807, 0.1_624, 0.1_016, -0.1_896, 0.0_263, 0.0_677, 0.2_310]],
			        [17, 0.55, [0.1_164, -0.0_216, 0.0_170, 0.1_589, -0.3_120, 0.1_005, -0.0_581, -0.1_458]],
			        [8, 0.89, [-0.1_758, -0.0_169, 0.1_004, -0.1_411, 0.1_312, 0.1_103, -0.1_996, 0.2_139]],
			        [3, 1000, [0.1_214, 0.0_352, -0.0_731, -0.1_562, -0.0_994, -0.0_906, -0.2_340, -0.0_539]],
			        # fmt: on
			    ]    )
			def 			lowerCamelCase  (self					,							__magic_name__					,							__magic_name__					,							__magic_name__    )     ->   Optional[int]:


							'''simple docstring'''



							snake_case_						:    List[str]				=  self.get_unet_model(model_id='''stabilityai/stable-diffusion-2'''					,							fpaa=SCREAMING_SNAKE_CASE_    )
							snake_case_						:    int				=  self.get_latents(SCREAMING_SNAKE_CASE_					,							shape=(4, 4, 96, 96)					,							fpaa=SCREAMING_SNAKE_CASE_    )
							snake_case_						:    Optional[int]				=  self.get_encoder_hidden_states(SCREAMING_SNAKE_CASE_					,							shape=(4, 77, 1024)					,							fpaa=SCREAMING_SNAKE_CASE_    )

							snake_case_						:    List[str]				=  model.apply(
							    {'''params''': params}					,							SCREAMING_SNAKE_CASE_					,							jnp.array(SCREAMING_SNAKE_CASE_					,							dtype=jnp.intaa    )					,							encoder_hidden_states=SCREAMING_SNAKE_CASE_					,							).sample

							assert sample.shape == latents.shape

							snake_case_						:    Optional[int]				=  jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())    )					,							dtype=jnp.floataa    )
							snake_case_						:    Tuple				=  jnp.array(SCREAMING_SNAKE_CASE_					,							dtype=jnp.floataa    )

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




'''simple docstring'''



import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()





class    UpperCAmelCase_  (_UpperCAmelCase					,		_UpperCAmelCase					,		unittest.TestCase  ):





    """simple docstring"""






    lowerCamelCase  :					Union[str, Any]							     =			VQModel
    lowerCamelCase  :					Union[str, Any]							     =			'sample'







    @property
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=(32, 32)	)				->     Any:
         __lowerCamelCase   :	Tuple								=		4
         __lowerCamelCase   :	Union[str, Any]								=		3

         __lowerCamelCase   :	List[Any]								=		floats_tensor((batch_size, num_channels) + sizes	).to(SCREAMING_SNAKE_CASE_	)

         return {"sample": image}







    @property
    def       lowercase_				(							self	)				->     Optional[int]:
         return (3, 32, 32)







    @property
    def       lowercase_				(							self	)				->     List[Any]:
         return (3, 32, 32)







    def       lowercase_				(							self	)				->     Optional[int]:
         __lowerCamelCase   :	Optional[Any]								=		{
             'block_out_channels': [32, 64],
             'in_channels': 3,
             'out_channels': 3,
             'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
             'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
             'latent_channels': 3,
         }
         __lowerCamelCase   :	Union[str, Any]								=		self.dummy_input
         return init_dict, inputs_dict







    def       lowercase_				(							self	)				->     str:
         pass







    def       lowercase_				(							self	)				->     Optional[int]:
         pass







    def       lowercase_				(							self	)				->     Any:
         __lowerCamelCase				,       __lowerCamelCase   :	Union[str, Any]								=		VQModel.from_pretrained('fusing/vqgan-dummy'			,     output_loading_info=SCREAMING_SNAKE_CASE_	)
         self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)
         self.assertEqual(len(loading_info['missing_keys']	)			,     0	)

         model.to(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		model(**self.dummy_input	)

         assert image is not None, "Make sure output is not None"







    def       lowercase_				(							self	)				->     int:
         __lowerCamelCase   :	List[Any]								=		VQModel.from_pretrained('fusing/vqgan-dummy'	)
         model.to(SCREAMING_SNAKE_CASE_	).eval()

         torch.manual_seed(0	)
         if torch.cuda.is_available():
              torch.cuda.manual_seed_all(0	)

         __lowerCamelCase   :	List[str]								=		torch.randn(1			,     model.config.in_channels			,     model.config.sample_size			,     model.config.sample_size	)
         __lowerCamelCase   :	Optional[int]								=		image.to(SCREAMING_SNAKE_CASE_	)
         with torch.no_grad():
              __lowerCamelCase   :	int								=		model(SCREAMING_SNAKE_CASE_	).sample

         __lowerCamelCase   :	List[str]								=		output[0, -1, -3:, -3:].flatten().cpu()
         # fmt: off
         __lowerCamelCase   :	Union[str, Any]								=		torch.tensor([-0.0_1_5_3, -0.4_0_4_4, -0.1_8_8_0, -0.5_1_6_1, -0.2_4_1_8, -0.4_0_7_2, -0.1_6_1_2, -0.0_6_3_3, -0.0_1_4_3]	)
         # fmt: on
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     atol=1E-3	)	)






