

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






     def 	UpperCAmelCase_				(				self		,						_lowercase		,						_lowercase	)->     Optional[Any]:
         return F"gaussian_noise_s={seed}_shape={'_'.join([str(_lowercase	) for s in shape]	)}.npy"






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






     def 	UpperCAmelCase_				(				self		,						_lowercase=0		,						_lowercase=(4, 4, 64, 64)		,						_lowercase=False	)->     Optional[Any]:
         UpperCamelCase_			  =						jnp.bfloataa if fpaa else jnp.floataa
         UpperCamelCase_			  =						jnp.array(load_hf_numpy(self.get_file_format(_lowercase		,						_lowercase	)	)		,						dtype=_lowercase	)
         return image






     def 	UpperCAmelCase_				(				self		,						_lowercase=False		,						_lowercase="CompVis/stable-diffusion-v1-4"	)->     Optional[int]:
         UpperCamelCase_			  =						jnp.bfloataa if fpaa else jnp.floataa
         UpperCamelCase_			  =						"bf16" if fpaa else None

         UpperCamelCase_  ,       UpperCamelCase_			  =						FlaxUNetaDConditionModel.from_pretrained(
             _lowercase		,						subfolder="unet"		,						dtype=_lowercase		,						revision=_lowercase	)
         return model, params






     def 	UpperCAmelCase_				(				self		,						_lowercase=0		,						_lowercase=(4, 77, 768)		,						_lowercase=False	)->     Optional[int]:
         UpperCamelCase_			  =						jnp.bfloataa if fpaa else jnp.floataa
         UpperCamelCase_			  =						jnp.array(load_hf_numpy(self.get_file_format(_lowercase		,						_lowercase	)	)		,						dtype=_lowercase	)
         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, 1_000, [-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 	UpperCAmelCase_				(				self		,						_lowercase		,						_lowercase		,						_lowercase	)->     List[Any]:
         UpperCamelCase_  ,       UpperCamelCase_			  =						self.get_unet_model(model_id="CompVis/stable-diffusion-v1-4"		,						fpaa=_lowercase	)
         UpperCamelCase_			  =						self.get_latents(_lowercase		,						fpaa=_lowercase	)
         UpperCamelCase_			  =						self.get_encoder_hidden_states(_lowercase		,						fpaa=_lowercase	)

         UpperCamelCase_			  =						model.apply(
             {"params": params}		,						_lowercase		,						jnp.array(_lowercase		,						dtype=jnp.intaa	)		,						encoder_hidden_states=_lowercase		,						).sample

         assert sample.shape == latents.shape

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

         # Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, in the same hardware
         assert jnp.allclose(_lowercase		,						_lowercase		,						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, 1_000, [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 	UpperCAmelCase_				(				self		,						_lowercase		,						_lowercase		,						_lowercase	)->     Optional[Any]:
         UpperCamelCase_  ,       UpperCamelCase_			  =						self.get_unet_model(model_id="stabilityai/stable-diffusion-2"		,						fpaa=_lowercase	)
         UpperCamelCase_			  =						self.get_latents(_lowercase		,						shape=(4, 4, 96, 96)		,						fpaa=_lowercase	)
         UpperCamelCase_			  =						self.get_encoder_hidden_states(_lowercase		,						shape=(4, 77, 1_024)		,						fpaa=_lowercase	)

         UpperCamelCase_			  =						model.apply(
             {"params": params}		,						_lowercase		,						jnp.array(_lowercase		,						dtype=jnp.intaa	)		,						encoder_hidden_states=_lowercase		,						).sample

         assert sample.shape == latents.shape

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

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



import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest






class 						__magic_name__					(							snake_case				):
     UpperCamelCase_	:Dict										=      (KDPMaDiscreteScheduler,)
     UpperCamelCase_	:str										=      1_0






     def 	UpperCAmelCase_				(				self		,						**_lowercase	)->     str:
         UpperCamelCase_			  =						{
             "num_train_timesteps": 1_100,
             "beta_start": 0.0_001,
             "beta_end": 0.02,
             "beta_schedule": "linear",
         }

         config.update(**_lowercase	)
         return config






     def 	UpperCAmelCase_				(				self	)->     Union[str, Any]:
         for timesteps in [10, 50, 100, 1_000]:
             self.check_over_configs(num_train_timesteps=_lowercase	)






     def 	UpperCAmelCase_				(				self	)->     int:
         for beta_start, beta_end in zip([0.00_001, 0.0_001, 0.001]		,						[0.0_002, 0.002, 0.02]	):
             self.check_over_configs(beta_start=_lowercase		,						beta_end=_lowercase	)






     def 	UpperCAmelCase_				(				self	)->     str:
         for schedule in ["linear", "scaled_linear"]:
             self.check_over_configs(beta_schedule=_lowercase	)






     def 	UpperCAmelCase_				(				self	)->     Any:
         for prediction_type in ["epsilon", "v_prediction"]:
             self.check_over_configs(prediction_type=_lowercase	)






     def 	UpperCAmelCase_				(				self	)->     Optional[int]:
         UpperCamelCase_			  =						self.scheduler_classes[0]
         UpperCamelCase_			  =						self.get_scheduler_config(prediction_type="v_prediction"	)
         UpperCamelCase_			  =						scheduler_class(**_lowercase	)

         scheduler.set_timesteps(self.num_inference_steps	)

         UpperCamelCase_			  =						self.dummy_model()
         UpperCamelCase_			  =						self.dummy_sample_deter * scheduler.init_noise_sigma
         UpperCamelCase_			  =						sample.to(_lowercase	)

         for i, t in enumerate(scheduler.timesteps	):
             UpperCamelCase_			  =						scheduler.scale_model_input(_lowercase		,						_lowercase	)

             UpperCamelCase_			  =						model(_lowercase		,						_lowercase	)

             UpperCamelCase_			  =						scheduler.step(_lowercase		,						_lowercase		,						_lowercase	)
             UpperCamelCase_			  =						output.prev_sample

         UpperCamelCase_			  =						torch.sum(torch.abs(_lowercase	)	)
         UpperCamelCase_			  =						torch.mean(torch.abs(_lowercase	)	)

         if torch_device in ["cpu", "mps"]:
             assert abs(result_sum.item() - 4.6_9_3_4e-0_7	) < 1e-2
             assert abs(result_mean.item() - 6.1_1_1_2e-1_0	) < 1e-3
         else:
             # CUDA
             assert abs(result_sum.item() - 4.6_9_3_4_2_8_6_5_0_1_7_0_9_7_2e-0_7	) < 1e-2
             assert abs(result_mean.item() - 0.0_002	) < 1e-3






     def 	UpperCAmelCase_				(				self	)->     Dict:
         if torch_device == "mps":
             return
         UpperCamelCase_			  =						self.scheduler_classes[0]
         UpperCamelCase_			  =						self.get_scheduler_config()
         UpperCamelCase_			  =						scheduler_class(**_lowercase	)

         scheduler.set_timesteps(self.num_inference_steps	)

         UpperCamelCase_			  =						self.dummy_model()
         UpperCamelCase_			  =						self.dummy_sample_deter * scheduler.init_noise_sigma
         UpperCamelCase_			  =						sample.to(_lowercase	)

         for i, t in enumerate(scheduler.timesteps	):
             UpperCamelCase_			  =						scheduler.scale_model_input(_lowercase		,						_lowercase	)

             UpperCamelCase_			  =						model(_lowercase		,						_lowercase	)

             UpperCamelCase_			  =						scheduler.step(_lowercase		,						_lowercase		,						_lowercase	)
             UpperCamelCase_			  =						output.prev_sample

         UpperCamelCase_			  =						torch.sum(torch.abs(_lowercase	)	)
         UpperCamelCase_			  =						torch.mean(torch.abs(_lowercase	)	)

         if torch_device in ["cpu", "mps"]:
             assert abs(result_sum.item() - 20.4_125	) < 1e-2
             assert abs(result_mean.item() - 0.0_266	) < 1e-3
         else:
             # CUDA
             assert abs(result_sum.item() - 20.4_125	) < 1e-2
             assert abs(result_mean.item() - 0.0_266	) < 1e-3






     def 	UpperCAmelCase_				(				self	)->     Optional[int]:
         if torch_device == "mps":
             return
         UpperCamelCase_			  =						self.scheduler_classes[0]
         UpperCamelCase_			  =						self.get_scheduler_config()
         UpperCamelCase_			  =						scheduler_class(**_lowercase	)

         scheduler.set_timesteps(self.num_inference_steps		,						device=_lowercase	)

         UpperCamelCase_			  =						self.dummy_model()
         UpperCamelCase_			  =						self.dummy_sample_deter.to(_lowercase	) * scheduler.init_noise_sigma

         for t in scheduler.timesteps:
             UpperCamelCase_			  =						scheduler.scale_model_input(_lowercase		,						_lowercase	)

             UpperCamelCase_			  =						model(_lowercase		,						_lowercase	)

             UpperCamelCase_			  =						scheduler.step(_lowercase		,						_lowercase		,						_lowercase	)
             UpperCamelCase_			  =						output.prev_sample

         UpperCamelCase_			  =						torch.sum(torch.abs(_lowercase	)	)
         UpperCamelCase_			  =						torch.mean(torch.abs(_lowercase	)	)

         if str(_lowercase	).startswith("cpu"	):
             # The following sum varies between 148 and 156 on mps. Why?
             assert abs(result_sum.item() - 20.4_125	) < 1e-2
             assert abs(result_mean.item() - 0.0_266	) < 1e-3
         else:
             # CUDA
             assert abs(result_sum.item() - 20.4_125	) < 1e-2
             assert abs(result_mean.item() - 0.0_266	) < 1e-3
