

import logging
import os
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


a_     =			{
    """debug""": logging.DEBUG,
    """info""": logging.INFO,
    """warning""": logging.WARNING,
    """error""": logging.ERROR,
    """critical""": logging.CRITICAL,
}

a_     =			logging.WARNING

def        lowerCamelCase__       (    ):
	SCREAMING_SNAKE_CASE      :					str    =				os.getenv("DATASETS_VERBOSITY"							,	UpperCAmelCase_)
	if env_level_str:
		if env_level_str in log_levels:
			return log_levels[env_level_str]
		else:
			logging.getLogger().warning(
			    f"Unknown option DATASETS_VERBOSITY={env_level_str}, "
			    f"has to be one of: { ', '.join(log_levels.keys()) }")
	return _default_log_level

def        lowerCamelCase__       (    ):
	return __name__.split(".")[0]

def        lowerCamelCase__       (    ):
	return logging.getLogger(_get_library_name())

def        lowerCamelCase__       (    ):
	# Apply our default configuration to the library root logger.
	SCREAMING_SNAKE_CASE      :					List[Any]    =				_get_library_root_logger()
	library_root_logger.setLevel(_get_default_logging_level())

def        lowerCamelCase__       (    ):
	SCREAMING_SNAKE_CASE      :					Any    =				_get_library_root_logger()
	library_root_logger.setLevel(logging.NOTSET)

def        lowerCamelCase__       (    _a = None):
	if name is None:
		SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				_get_library_name()
	return logging.getLogger(UpperCAmelCase_)

def        lowerCamelCase__       (    ):
	return _get_library_root_logger().getEffectiveLevel()

def        lowerCamelCase__       (    _a):
	_get_library_root_logger().setLevel(UpperCAmelCase_)

def        lowerCamelCase__       (    ):
	return set_verbosity(UpperCAmelCase_)

def        lowerCamelCase__       (    ):
	return set_verbosity(UpperCAmelCase_)

def        lowerCamelCase__       (    ):
	return set_verbosity(UpperCAmelCase_)

def        lowerCamelCase__       (    ):
	return set_verbosity(UpperCAmelCase_)

def        lowerCamelCase__       (    ):
	SCREAMING_SNAKE_CASE      :					List[str]    =				False

def        lowerCamelCase__       (    ):
	SCREAMING_SNAKE_CASE      :					List[Any]    =				True


# Configure the library root logger at the module level (singleton-like)
_configure_library_root_logger()






class _UpperCamelCase							:
			'''simple docstring'''







			def __init__(	self							:       List[Any]					,					*a							:       List[str]					,					**a							:       Any		)							->				Any:  # pylint: disable=unused-argument

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					str    =				args[0] if args else None
			def __iter__(	self							:       Any		)							->				str:

				"""simple docstring"""

				return iter(self._iterator		)
			def __getattr__(	self							:       Optional[Any]					,					a							:       Any		)							->				Dict:

				"""simple docstring"""

				def empty_fn(*a							:       Optional[Any]					,					**a							:       Optional[Any]		):  # pylint: disable=unused-argument
					return

				return empty_fn
			def __enter__(	self							:       int		)							->				List[Any]:

				"""simple docstring"""

				return self






			def __exit__(	self							:       Optional[int]					,					a							:       List[Any]					,					a							:       Any					,					a							:       Union[str, Any]		)							->				List[str]:

				"""simple docstring"""

				return

a_     =			True






class _UpperCamelCase							:
			'''simple docstring'''







			def __call__(	self							:       str					,					*a							:       int					,					a							:       str=False					,					**a							:       int		)							->				str:

				"""simple docstring"""

				if _tqdm_active and not disable:
					return tqdm_lib.tqdm(*SCREAMING_SNAKE_CASE_					,					**SCREAMING_SNAKE_CASE_		)
				else:
					return EmptyTqdm(*SCREAMING_SNAKE_CASE_					,					**SCREAMING_SNAKE_CASE_		)
			def 							__UpperCamelCase       (	self							:       int					,					*a							:       int					,					**a							:       List[str]		)							->				Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE      :					Union[str, Any]    =				None
				if _tqdm_active:
					return tqdm_lib.tqdm.set_lock(*SCREAMING_SNAKE_CASE_					,					**SCREAMING_SNAKE_CASE_		)






			def 							__UpperCamelCase       (	self							:       Optional[Any]		)							->				Union[str, Any]:

				"""simple docstring"""

				if _tqdm_active:
					return tqdm_lib.tqdm.get_lock()

a_     =			_tqdm_cls()

def        lowerCamelCase__       (    ):
	global _tqdm_active
	return bool(_tqdm_active)

def        lowerCamelCase__       (    ):
	global _tqdm_active
	SCREAMING_SNAKE_CASE      :					Optional[int]    =				True


def        lowerCamelCase__       (    ):
	global _tqdm_active
	SCREAMING_SNAKE_CASE      :					Tuple    =				False



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





    """simple docstring"""






    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[int]:
         return f'gaussian_noise_s={seed}_shape={"_".join([str(SCREAMING_SNAKE_CASE_	) for s in shape]	)}.npy'







    def       lowercase_				(							self	)				->     Optional[Any]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=0			,     SCREAMING_SNAKE_CASE_=(4, 4, 64, 64)			,     SCREAMING_SNAKE_CASE_=False	)				->     str:
         __lowerCamelCase   :	List[str]								=		jnp.bfloataa if fpaa else jnp.floataa
         __lowerCamelCase   :	Tuple								=		jnp.array(load_hf_numpy(self.get_file_format(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)	)			,     dtype=SCREAMING_SNAKE_CASE_	)
         return image







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=False			,     SCREAMING_SNAKE_CASE_="CompVis/stable-diffusion-v1-4"	)				->     Dict:
         __lowerCamelCase   :	Union[str, Any]								=		jnp.bfloataa if fpaa else jnp.floataa
         __lowerCamelCase   :	Optional[Any]								=		'bf16' if fpaa else None

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







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=0			,     SCREAMING_SNAKE_CASE_=(4, 77, 7_68)			,     SCREAMING_SNAKE_CASE_=False	)				->     Union[str, Any]:
         __lowerCamelCase   :	Any								=		jnp.bfloataa if fpaa else jnp.floataa
         __lowerCamelCase   :	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_3_2_3, -0.1_3_0_4, 0.0_8_1_3, -0.3_0_9_3, -0.0_9_1_9, -0.1_5_7_1, -0.1_1_2_5, -0.5_8_0_6]],
            [17, 0.5_5, [-0.0_8_3_1, -0.2_4_4_3, 0.0_9_0_1, -0.0_9_1_9, 0.3_3_9_6, 0.0_1_0_3, -0.3_7_4_3, 0.0_7_0_1]],
            [8, 0.8_9, [-0.4_8_6_3, 0.0_8_5_9, 0.0_8_7_5, -0.1_6_5_8, 0.9_1_9_9, -0.0_1_1_4, 0.4_8_3_9, 0.4_6_3_9]],
            [3, 10_00, [-0.5_6_4_9, 0.2_4_0_2, -0.5_5_1_8, 0.1_2_4_8, 1.1_3_2_8, -0.2_4_4_3, -0.0_3_2_5, -1.0_0_7_8]],
            # fmt: on
        ]	)
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase				,       __lowerCamelCase   :	Tuple								=		self.get_unet_model(model_id='CompVis/stable-diffusion-v1-4'			,     fpaa=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		self.get_latents(SCREAMING_SNAKE_CASE_			,     fpaa=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	int								=		self.get_encoder_hidden_states(SCREAMING_SNAKE_CASE_			,     fpaa=SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	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

         __lowerCamelCase   :	Any								=		jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())	)			,     dtype=jnp.floataa	)
         __lowerCamelCase   :	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_5_1_4, 0.0_8_0_7, 0.1_6_2_4, 0.1_0_1_6, -0.1_8_9_6, 0.0_2_6_3, 0.0_6_7_7, 0.2_3_1_0]],
            [17, 0.5_5, [0.1_1_6_4, -0.0_2_1_6, 0.0_1_7_0, 0.1_5_8_9, -0.3_1_2_0, 0.1_0_0_5, -0.0_5_8_1, -0.1_4_5_8]],
            [8, 0.8_9, [-0.1_7_5_8, -0.0_1_6_9, 0.1_0_0_4, -0.1_4_1_1, 0.1_3_1_2, 0.1_1_0_3, -0.1_9_9_6, 0.2_1_3_9]],
            [3, 10_00, [0.1_2_1_4, 0.0_3_5_2, -0.0_7_3_1, -0.1_5_6_2, -0.0_9_9_4, -0.0_9_0_6, -0.2_3_4_0, -0.0_5_3_9]],
            # fmt: on
        ]	)
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[int]:
         __lowerCamelCase				,       __lowerCamelCase   :	List[str]								=		self.get_unet_model(model_id='stabilityai/stable-diffusion-2'			,     fpaa=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	int								=		self.get_latents(SCREAMING_SNAKE_CASE_			,     shape=(4, 4, 96, 96)			,     fpaa=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[int]								=		self.get_encoder_hidden_states(SCREAMING_SNAKE_CASE_			,     shape=(4, 77, 10_24)			,     fpaa=SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	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

         __lowerCamelCase   :	Optional[int]								=		jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())	)			,     dtype=jnp.floataa	)
         __lowerCamelCase   :	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	)






