
import torch

from diffusers import EulerDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest





class    __A(					_lowerCamelCase							):
					SCREAMING_SNAKE_CASE           =     (EulerDiscreteScheduler,)
					SCREAMING_SNAKE_CASE           =     1_0

					def  lowercase__   (						self       :  Union[str, Any] ,       **__UpperCamelCase       :  Optional[Any] ):
												lowerCamelCase_					=       {
												    """num_train_timesteps""": 1_1_0_0,
												    """beta_start""": 0.0001,
												    """beta_end""": 0.02,
												    """beta_schedule""": """linear""",
												}

												config.update(**__UpperCamelCase )
												return config

					def  lowercase__   (						self       :  int ):
												for timesteps in [1_0, 5_0, 1_0_0, 1_0_0_0]:
																			self.check_over_configs(num_train_timesteps=__UpperCamelCase )

					def  lowercase__   (						self       :  Union[str, Any] ):
												for beta_start, beta_end in zip([0.00001, 0.0001, 0.001] ,       [0.0002, 0.002, 0.02] ):
																			self.check_over_configs(beta_start=__UpperCamelCase ,       beta_end=__UpperCamelCase )

					def  lowercase__   (						self       :  Optional[int] ):
												for schedule in ["linear", "scaled_linear"]:
																			self.check_over_configs(beta_schedule=__UpperCamelCase )

					def  lowercase__   (						self       :  int ):
												for prediction_type in ["epsilon", "v_prediction"]:
																			self.check_over_configs(prediction_type=__UpperCamelCase )

					def  lowercase__   (						self       :  int ):
												lowerCamelCase_					=       self.scheduler_classes[0]
												lowerCamelCase_					=       self.get_scheduler_config()
												lowerCamelCase_					=       scheduler_class(**__UpperCamelCase )

												scheduler.set_timesteps(self.num_inference_steps )

												lowerCamelCase_					=       torch.manual_seed(0 )

												lowerCamelCase_					=       self.dummy_model()
												lowerCamelCase_					=       self.dummy_sample_deter * scheduler.init_noise_sigma
												lowerCamelCase_					=       sample.to(__UpperCamelCase )

												for i, t in enumerate(scheduler.timesteps ):
																			lowerCamelCase_					=       scheduler.scale_model_input(__UpperCamelCase ,       __UpperCamelCase )

																			lowerCamelCase_					=       model(__UpperCamelCase ,       __UpperCamelCase )

																			lowerCamelCase_					=       scheduler.step(__UpperCamelCase ,       __UpperCamelCase ,       __UpperCamelCase ,       generator=__UpperCamelCase )
																			lowerCamelCase_					=       output.prev_sample

												lowerCamelCase_					=       torch.sum(torch.abs(__UpperCamelCase ) )
												lowerCamelCase_					=       torch.mean(torch.abs(__UpperCamelCase ) )

												assert abs(result_sum.item() - 1_0.0_8_0_7 ) < 1E-2
												assert abs(result_mean.item() - 0.0131 ) < 1E-3

					def  lowercase__   (						self       :  List[Any] ):
												lowerCamelCase_					=       self.scheduler_classes[0]
												lowerCamelCase_					=       self.get_scheduler_config(prediction_type="""v_prediction""" )
												lowerCamelCase_					=       scheduler_class(**__UpperCamelCase )

												scheduler.set_timesteps(self.num_inference_steps )

												lowerCamelCase_					=       torch.manual_seed(0 )

												lowerCamelCase_					=       self.dummy_model()
												lowerCamelCase_					=       self.dummy_sample_deter * scheduler.init_noise_sigma
												lowerCamelCase_					=       sample.to(__UpperCamelCase )

												for i, t in enumerate(scheduler.timesteps ):
																			lowerCamelCase_					=       scheduler.scale_model_input(__UpperCamelCase ,       __UpperCamelCase )

																			lowerCamelCase_					=       model(__UpperCamelCase ,       __UpperCamelCase )

																			lowerCamelCase_					=       scheduler.step(__UpperCamelCase ,       __UpperCamelCase ,       __UpperCamelCase ,       generator=__UpperCamelCase )
																			lowerCamelCase_					=       output.prev_sample

												lowerCamelCase_					=       torch.sum(torch.abs(__UpperCamelCase ) )
												lowerCamelCase_					=       torch.mean(torch.abs(__UpperCamelCase ) )

												assert abs(result_sum.item() - 0.0002 ) < 1E-2
												assert abs(result_mean.item() - 2.2676E-06 ) < 1E-3

					def  lowercase__   (						self       :  Union[str, Any] ):
												lowerCamelCase_					=       self.scheduler_classes[0]
												lowerCamelCase_					=       self.get_scheduler_config()
												lowerCamelCase_					=       scheduler_class(**__UpperCamelCase )

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

												lowerCamelCase_					=       torch.manual_seed(0 )

												lowerCamelCase_					=       self.dummy_model()
												lowerCamelCase_					=       self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
												lowerCamelCase_					=       sample.to(__UpperCamelCase )

												for t in scheduler.timesteps:
																			lowerCamelCase_					=       scheduler.scale_model_input(__UpperCamelCase ,       __UpperCamelCase )

																			lowerCamelCase_					=       model(__UpperCamelCase ,       __UpperCamelCase )

																			lowerCamelCase_					=       scheduler.step(__UpperCamelCase ,       __UpperCamelCase ,       __UpperCamelCase ,       generator=__UpperCamelCase )
																			lowerCamelCase_					=       output.prev_sample

												lowerCamelCase_					=       torch.sum(torch.abs(__UpperCamelCase ) )
												lowerCamelCase_					=       torch.mean(torch.abs(__UpperCamelCase ) )

												assert abs(result_sum.item() - 1_0.0_8_0_7 ) < 1E-2
												assert abs(result_mean.item() - 0.0131 ) < 1E-3



					def  lowercase__   (						self       :  Optional[Any] ):
												lowerCamelCase_					=       self.scheduler_classes[0]
												lowerCamelCase_					=       self.get_scheduler_config()
												lowerCamelCase_					=       scheduler_class(**__UpperCamelCase ,       use_karras_sigmas=__UpperCamelCase )

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

												lowerCamelCase_					=       torch.manual_seed(0 )

												lowerCamelCase_					=       self.dummy_model()
												lowerCamelCase_					=       self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
												lowerCamelCase_					=       sample.to(__UpperCamelCase )

												for t in scheduler.timesteps:
																			lowerCamelCase_					=       scheduler.scale_model_input(__UpperCamelCase ,       __UpperCamelCase )

																			lowerCamelCase_					=       model(__UpperCamelCase ,       __UpperCamelCase )

																			lowerCamelCase_					=       scheduler.step(__UpperCamelCase ,       __UpperCamelCase ,       __UpperCamelCase ,       generator=__UpperCamelCase )
																			lowerCamelCase_					=       output.prev_sample

												lowerCamelCase_					=       torch.sum(torch.abs(__UpperCamelCase ) )
												lowerCamelCase_					=       torch.mean(torch.abs(__UpperCamelCase ) )

												assert abs(result_sum.item() - 1_2_4.5_2_2_9_9_4_9_9_5_1_1_7_1_9 ) < 1E-2
												assert abs(result_mean.item() - 0.16213932633399963 ) < 1E-3










import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible



class        lowerCAmelCase__					(  unittest.TestCase       ):


				def __UpperCamelCase						(							self      :  List[str]       )		->     Any:
											A      						=    [
											    'safety_checker/pytorch_model.bin',
											    'safety_checker/model.safetensors',
											    'vae/diffusion_pytorch_model.bin',
											    'vae/diffusion_pytorch_model.safetensors',
											    'text_encoder/pytorch_model.bin',
											    'text_encoder/model.safetensors',
											    'unet/diffusion_pytorch_model.bin',
											    'unet/diffusion_pytorch_model.safetensors',
											]
											self.assertTrue(is_safetensors_compatible(__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  str       )		->     List[Any]:
											A      						=    [
											    'unet/diffusion_pytorch_model.bin',
											    'unet/diffusion_pytorch_model.safetensors',
											]
											self.assertTrue(is_safetensors_compatible(__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  Any       )		->     Tuple:
											A      						=    [
											    'safety_checker/pytorch_model.bin',
											    'safety_checker/model.safetensors',
											    'vae/diffusion_pytorch_model.bin',
											    'vae/diffusion_pytorch_model.safetensors',
											    'text_encoder/pytorch_model.bin',
											    'text_encoder/model.safetensors',
											    'unet/diffusion_pytorch_model.bin',
											    # Removed: 'unet/diffusion_pytorch_model.safetensors',
											]
											self.assertFalse(is_safetensors_compatible(__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  Dict       )		->     Any:
											A      						=    [
											    'text_encoder/pytorch_model.bin',
											    'text_encoder/model.safetensors',
											]
											self.assertTrue(is_safetensors_compatible(__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  Optional[int]       )		->     List[str]:
											A      						=    [
											    'safety_checker/pytorch_model.bin',
											    'safety_checker/model.safetensors',
											    'vae/diffusion_pytorch_model.bin',
											    'vae/diffusion_pytorch_model.safetensors',
											    'text_encoder/pytorch_model.bin',
											    # Removed: 'text_encoder/model.safetensors',
											    'unet/diffusion_pytorch_model.bin',
											    'unet/diffusion_pytorch_model.safetensors',
											]
											self.assertFalse(is_safetensors_compatible(__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  List[str]       )		->     List[str]:
											A      						=    [
											    'safety_checker/pytorch_model.fp16.bin',
											    'safety_checker/model.fp16.safetensors',
											    'vae/diffusion_pytorch_model.fp16.bin',
											    'vae/diffusion_pytorch_model.fp16.safetensors',
											    'text_encoder/pytorch_model.fp16.bin',
											    'text_encoder/model.fp16.safetensors',
											    'unet/diffusion_pytorch_model.fp16.bin',
											    'unet/diffusion_pytorch_model.fp16.safetensors',
											]
											A      						=    'fp16'
											self.assertTrue(is_safetensors_compatible(__UpperCamelCase						, variant=__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  Optional[Any]       )		->     Optional[Any]:
											A      						=    [
											    'unet/diffusion_pytorch_model.fp16.bin',
											    'unet/diffusion_pytorch_model.fp16.safetensors',
											]
											A      						=    'fp16'
											self.assertTrue(is_safetensors_compatible(__UpperCamelCase						, variant=__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  int       )		->     Optional[int]:
											# pass variant but use the non-variant filenames
											A      						=    [
											    'unet/diffusion_pytorch_model.bin',
											    'unet/diffusion_pytorch_model.safetensors',
											]
											A      						=    'fp16'
											self.assertTrue(is_safetensors_compatible(__UpperCamelCase						, variant=__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  Any       )		->     List[str]:
											A      						=    [
											    'safety_checker/pytorch_model.fp16.bin',
											    'safety_checker/model.fp16.safetensors',
											    'vae/diffusion_pytorch_model.fp16.bin',
											    'vae/diffusion_pytorch_model.fp16.safetensors',
											    'text_encoder/pytorch_model.fp16.bin',
											    'text_encoder/model.fp16.safetensors',
											    'unet/diffusion_pytorch_model.fp16.bin',
											    # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
											]
											A      						=    'fp16'
											self.assertFalse(is_safetensors_compatible(__UpperCamelCase						, variant=__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  Optional[Any]       )		->     Optional[int]:
											A      						=    [
											    'text_encoder/pytorch_model.fp16.bin',
											    'text_encoder/model.fp16.safetensors',
											]
											A      						=    'fp16'
											self.assertTrue(is_safetensors_compatible(__UpperCamelCase						, variant=__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  List[str]       )		->     int:
											# pass variant but use the non-variant filenames
											A      						=    [
											    'text_encoder/pytorch_model.bin',
											    'text_encoder/model.safetensors',
											]
											A      						=    'fp16'
											self.assertTrue(is_safetensors_compatible(__UpperCamelCase						, variant=__UpperCamelCase       )       )


				def __UpperCamelCase						(							self      :  List[str]       )		->     Optional[Any]:
											A      						=    [
											    'safety_checker/pytorch_model.fp16.bin',
											    'safety_checker/model.fp16.safetensors',
											    'vae/diffusion_pytorch_model.fp16.bin',
											    'vae/diffusion_pytorch_model.fp16.safetensors',
											    'text_encoder/pytorch_model.fp16.bin',
											    # 'text_encoder/model.fp16.safetensors',
											    'unet/diffusion_pytorch_model.fp16.bin',
											    'unet/diffusion_pytorch_model.fp16.safetensors',
											]
											A      						=    'fp16'
											self.assertFalse(is_safetensors_compatible(__UpperCamelCase						, variant=__UpperCamelCase       )       )