
"""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 ):

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







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

							def _lowercase       (self     ):

									"""simple docstring"""







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

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_=0							,							SCREAMING_SNAKE_CASE_=(4, 4, 64, 64)							,							SCREAMING_SNAKE_CASE_=False     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			jnp.bfloataa if fpaa else jnp.floataa
									SCREAMING_SNAKE_CASE_    			=			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"     ):

									"""simple docstring"""







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

									SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			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     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			jnp.bfloataa if fpaa else jnp.floataa
									SCREAMING_SNAKE_CASE_    			=			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.23_23, -0.13_04, 0.08_13, -0.30_93, -0.09_19, -0.15_71, -0.11_25, -0.58_06]],
							        [17, 0.55, [-0.08_31, -0.24_43, 0.09_01, -0.09_19, 0.33_96, 0.01_03, -0.37_43, 0.07_01]],
							        [8, 0.89, [-0.48_63, 0.08_59, 0.08_75, -0.16_58, 0.91_99, -0.01_14, 0.48_39, 0.46_39]],
							        [3, 10_00, [-0.56_49, 0.24_02, -0.55_18, 0.12_48, 1.13_28, -0.24_43, -0.03_25, -1.00_78]],
							        # fmt: on
							    ]     )
							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







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

									SCREAMING_SNAKE_CASE_    			=			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

									SCREAMING_SNAKE_CASE_    			=			jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())     )							,							dtype=jnp.floataa     )
									SCREAMING_SNAKE_CASE_    			=			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.15_14, 0.08_07, 0.16_24, 0.10_16, -0.18_96, 0.02_63, 0.06_77, 0.23_10]],
							        [17, 0.55, [0.11_64, -0.02_16, 0.01_70, 0.15_89, -0.31_20, 0.10_05, -0.05_81, -0.14_58]],
							        [8, 0.89, [-0.17_58, -0.01_69, 0.10_04, -0.14_11, 0.13_12, 0.11_03, -0.19_96, 0.21_39]],
							        [3, 10_00, [0.12_14, 0.03_52, -0.07_31, -0.15_62, -0.09_94, -0.09_06, -0.23_40, -0.05_39]],
							        # fmt: on
							    ]     )
							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			self.get_unet_model(model_id='''stabilityai/stable-diffusion-2'''							,							fpaa=SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			self.get_latents(SCREAMING_SNAKE_CASE_							,							shape=(4, 4, 96, 96)							,							fpaa=SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			self.get_encoder_hidden_states(SCREAMING_SNAKE_CASE_							,							shape=(4, 77, 10_24)							,							fpaa=SCREAMING_SNAKE_CASE_     )

									SCREAMING_SNAKE_CASE_    			=			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

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




from ..utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
)


try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						from ..utils.dummy_pt_objects import *  # noqa F403
else:
						from .scheduling_consistency_models import CMStochasticIterativeScheduler
						from .scheduling_ddim import DDIMScheduler
						from .scheduling_ddim_inverse import DDIMInverseScheduler
						from .scheduling_ddim_parallel import DDIMParallelScheduler
						from .scheduling_ddpm import DDPMScheduler
						from .scheduling_ddpm_parallel import DDPMParallelScheduler
						from .scheduling_deis_multistep import DEISMultistepScheduler
						from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
						from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler
						from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
						from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
						from .scheduling_euler_discrete import EulerDiscreteScheduler
						from .scheduling_heun_discrete import HeunDiscreteScheduler
						from .scheduling_ipndm import IPNDMScheduler
						from .scheduling_k_dpm_2_ancestral_discrete import KDPMaAncestralDiscreteScheduler
						from .scheduling_k_dpm_2_discrete import KDPMaDiscreteScheduler
						from .scheduling_karras_ve import KarrasVeScheduler
						from .scheduling_pndm import PNDMScheduler
						from .scheduling_repaint import RePaintScheduler
						from .scheduling_sde_ve import ScoreSdeVeScheduler
						from .scheduling_sde_vp import ScoreSdeVpScheduler
						from .scheduling_unclip import UnCLIPScheduler
						from .scheduling_unipc_multistep import UniPCMultistepScheduler
						from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
						from .scheduling_vq_diffusion import VQDiffusionScheduler

try:
						if not is_flax_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						from ..utils.dummy_flax_objects import *  # noqa F403
else:
						from .scheduling_ddim_flax import FlaxDDIMScheduler
						from .scheduling_ddpm_flax import FlaxDDPMScheduler
						from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler
						from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler
						from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler
						from .scheduling_pndm_flax import FlaxPNDMScheduler
						from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler
						from .scheduling_utils_flax import (
						    FlaxKarrasDiffusionSchedulers,
						    FlaxSchedulerMixin,
						    FlaxSchedulerOutput,
						    broadcast_to_shape_from_left,
						)


try:
						if not (is_torch_available() and is_scipy_available()):
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						from ..utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
						from .scheduling_lms_discrete import LMSDiscreteScheduler

try:
						if not (is_torch_available() and is_torchsde_available()):
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						from ..utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
						from .scheduling_dpmsolver_sde import DPMSolverSDEScheduler