





"""simple docstring"""




import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest


class  __UpperCamelCase       (		a__		):
							lowerCamelCase     :      Tuple         =(KDPMaDiscreteScheduler,)
							lowerCamelCase     :      Any         =10






							def 					__a     (     self   ,      **lowerCAmelCase__   )			->				Optional[Any]:
								a   :	Optional[Any]			=       {
								    "num_train_timesteps": 1100,
								    "beta_start": 0.0_001,
								    "beta_end": 0.02,
								    "beta_schedule": "linear",
								}

								config.update(**lowerCAmelCase__   )
								return config






							def 					__a     (     self   )			->				Optional[int]:
								for timesteps in [10, 50, 100, 1000]:
									self.check_over_configs(num_train_timesteps=lowerCAmelCase__   )






							def 					__a     (     self   )			->				Optional[Any]:
								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=lowerCAmelCase__   ,      beta_end=lowerCAmelCase__   )






							def 					__a     (     self   )			->				Dict:
								for schedule in ["linear", "scaled_linear"]:
									self.check_over_configs(beta_schedule=lowerCAmelCase__   )






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






							def 					__a     (     self   )			->				Dict:
								a   :	Optional[Any]			=       self.scheduler_classes[0]
								a   :	Dict			=       self.get_scheduler_config(prediction_type="v_prediction"   )
								a   :	Dict			=       scheduler_class(**lowerCAmelCase__   )

								scheduler.set_timesteps(self.num_inference_steps   )

								a   :	List[str]			=       self.dummy_model()
								a   :	Dict			=       self.dummy_sample_deter * scheduler.init_noise_sigma
								a   :	Dict			=       sample.to(lowerCAmelCase__   )

								for i, t in enumerate(scheduler.timesteps   ):
									a   :	List[str]			=       scheduler.scale_model_input(lowerCAmelCase__   ,      lowerCAmelCase__   )

									a   :	List[str]			=       model(lowerCAmelCase__   ,      lowerCAmelCase__   )

									a   :	List[str]			=       scheduler.step(lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )
									a   :	Optional[Any]			=       output.prev_sample

								a   :	Tuple			=       torch.sum(torch.abs(lowerCAmelCase__   )   )
								a   :	List[Any]			=       torch.mean(torch.abs(lowerCAmelCase__   )   )

								if torch_device in ["cpu", "mps"]:
									assert abs(result_sum.item() - 4.6_934E-07   ) < 1E-2
									assert abs(result_mean.item() - 6.1_112E-10   ) < 1E-3
								else:
									# CUDA
									assert abs(result_sum.item() - 4.693_428_650_170_972E-07   ) < 1E-2
									assert abs(result_mean.item() - 0.0_002   ) < 1E-3






							def 					__a     (     self   )			->				Tuple:
								if torch_device == "mps":
									return
								a   :	str			=       self.scheduler_classes[0]
								a   :	str			=       self.get_scheduler_config()
								a   :	Dict			=       scheduler_class(**lowerCAmelCase__   )

								scheduler.set_timesteps(self.num_inference_steps   )

								a   :	Dict			=       self.dummy_model()
								a   :	List[Any]			=       self.dummy_sample_deter * scheduler.init_noise_sigma
								a   :	Tuple			=       sample.to(lowerCAmelCase__   )

								for i, t in enumerate(scheduler.timesteps   ):
									a   :	List[str]			=       scheduler.scale_model_input(lowerCAmelCase__   ,      lowerCAmelCase__   )

									a   :	Union[str, Any]			=       model(lowerCAmelCase__   ,      lowerCAmelCase__   )

									a   :	List[Any]			=       scheduler.step(lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )
									a   :	int			=       output.prev_sample

								a   :	List[str]			=       torch.sum(torch.abs(lowerCAmelCase__   )   )
								a   :	List[str]			=       torch.mean(torch.abs(lowerCAmelCase__   )   )

								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 					__a     (     self   )			->				Dict:
								if torch_device == "mps":
									return
								a   :	Any			=       self.scheduler_classes[0]
								a   :	Tuple			=       self.get_scheduler_config()
								a   :	List[Any]			=       scheduler_class(**lowerCAmelCase__   )

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

								a   :	Optional[Any]			=       self.dummy_model()
								a   :	int			=       self.dummy_sample_deter.to(lowerCAmelCase__   ) * scheduler.init_noise_sigma

								for t in scheduler.timesteps:
									a   :	str			=       scheduler.scale_model_input(lowerCAmelCase__   ,      lowerCAmelCase__   )

									a   :	Optional[int]			=       model(lowerCAmelCase__   ,      lowerCAmelCase__   )

									a   :	List[str]			=       scheduler.step(lowerCAmelCase__   ,      lowerCAmelCase__   ,      lowerCAmelCase__   )
									a   :	Dict			=       output.prev_sample

								a   :	Dict			=       torch.sum(torch.abs(lowerCAmelCase__   )   )
								a   :	List[str]			=       torch.mean(torch.abs(lowerCAmelCase__   )   )

								if str(lowerCAmelCase__   ).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









"""simple docstring"""




import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible


class  __UpperCamelCase       (		unittest.TestCase		):






							def 					__a     (     self   )			->				Optional[Any]:
								a   :	Optional[int]			=       [
								    "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(lowerCAmelCase__   )   )






							def 					__a     (     self   )			->				Optional[Any]:
								a   :	str			=       [
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								self.assertTrue(is_safetensors_compatible(lowerCAmelCase__   )   )






							def 					__a     (     self   )			->				Dict:
								a   :	List[str]			=       [
								    "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(lowerCAmelCase__   )   )






							def 					__a     (     self   )			->				List[Any]:
								a   :	Optional[Any]			=       [
								    "text_encoder/pytorch_model.bin",
								    "text_encoder/model.safetensors",
								]
								self.assertTrue(is_safetensors_compatible(lowerCAmelCase__   )   )






							def 					__a     (     self   )			->				Tuple:
								a   :	Tuple			=       [
								    "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(lowerCAmelCase__   )   )






							def 					__a     (     self   )			->				Dict:
								a   :	Dict			=       [
								    "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   :	Dict			=       "fp16"
								self.assertTrue(is_safetensors_compatible(lowerCAmelCase__   ,      variant=lowerCAmelCase__   )   )






							def 					__a     (     self   )			->				List[str]:
								a   :	List[Any]			=       [
								    "unet/diffusion_pytorch_model.fp16.bin",
								    "unet/diffusion_pytorch_model.fp16.safetensors",
								]
								a   :	Any			=       "fp16"
								self.assertTrue(is_safetensors_compatible(lowerCAmelCase__   ,      variant=lowerCAmelCase__   )   )






							def 					__a     (     self   )			->				int:
								# pass variant but use the non-variant filenames
								a   :	int			=       [
								    "unet/diffusion_pytorch_model.bin",
								    "unet/diffusion_pytorch_model.safetensors",
								]
								a   :	Tuple			=       "fp16"
								self.assertTrue(is_safetensors_compatible(lowerCAmelCase__   ,      variant=lowerCAmelCase__   )   )






							def 					__a     (     self   )			->				str:
								a   :	str			=       [
								    "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   :	Any			=       "fp16"
								self.assertFalse(is_safetensors_compatible(lowerCAmelCase__   ,      variant=lowerCAmelCase__   )   )






							def 					__a     (     self   )			->				str:
								a   :	Union[str, Any]			=       [
								    "text_encoder/pytorch_model.fp16.bin",
								    "text_encoder/model.fp16.safetensors",
								]
								a   :	str			=       "fp16"
								self.assertTrue(is_safetensors_compatible(lowerCAmelCase__   ,      variant=lowerCAmelCase__   )   )






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






							def 					__a     (     self   )			->				Optional[Any]:
								a   :	Any			=       [
								    "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   :	Optional[int]			=       "fp16"
								self.assertFalse(is_safetensors_compatible(lowerCAmelCase__   ,      variant=lowerCAmelCase__   )   )


