



'''simple docstring'''





import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest







class        UpperCAmelCase_			(			_a      ):
    """simple docstring"""







    lowercase   =						(KDPMaDiscreteScheduler,)
    lowercase   =						10




    def lowerCamelCase			(       self   :	Optional[int]					,	**snake_case_   :	List[str]		):
          snake_case__    :       List[Any]       		= {
              """num_train_timesteps""": 1_100,
              """beta_start""": 0.0001,
              """beta_end""": 0.02,
              """beta_schedule""": """linear""",
          }

          config.update(**snake_case_		)
          return config




    def lowerCamelCase			(       self   :	Union[str, Any]		):
          for timesteps in [10, 50, 100, 1_000]:
                self.check_over_configs(num_train_timesteps=snake_case_		)




    def lowerCamelCase			(       self   :	List[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=snake_case_					,	beta_end=snake_case_		)




    def lowerCamelCase			(       self   :	List[Any]		):
          for schedule in ["linear", "scaled_linear"]:
                self.check_over_configs(beta_schedule=snake_case_		)




    def lowerCamelCase			(       self   :	Optional[int]		):
          for prediction_type in ["epsilon", "v_prediction"]:
                self.check_over_configs(prediction_type=snake_case_		)




    def lowerCamelCase			(       self   :	List[str]		):
          snake_case__    :       Any       		= self.scheduler_classes[0]
          snake_case__    :       List[str]       		= self.get_scheduler_config(prediction_type="""v_prediction"""		)
          snake_case__    :       Union[str, Any]       		= scheduler_class(**snake_case_		)

          scheduler.set_timesteps(self.num_inference_steps		)

          snake_case__    :       Tuple       		= self.dummy_model()
          snake_case__    :       Union[str, Any]       		= self.dummy_sample_deter * scheduler.init_noise_sigma
          snake_case__    :       List[Any]       		= sample.to(snake_case_		)

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

                snake_case__    :       Union[str, Any]       		= model(snake_case_					,	snake_case_		)

                snake_case__    :       Dict       		= scheduler.step(snake_case_					,	snake_case_					,	snake_case_		)
                snake_case__    :       Optional[int]       		= output.prev_sample

          snake_case__    :       Optional[Any]       		= torch.sum(torch.abs(snake_case_		)		)
          snake_case__    :       Any       		= torch.mean(torch.abs(snake_case_		)		)

          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.0002		) < 1E-3




    def lowerCamelCase			(       self   :	Optional[int]		):
          if torch_device == "mps":
                return
          snake_case__    :       int       		= self.scheduler_classes[0]
          snake_case__    :       Union[str, Any]       		= self.get_scheduler_config()
          snake_case__    :       List[str]       		= scheduler_class(**snake_case_		)

          scheduler.set_timesteps(self.num_inference_steps		)

          snake_case__    :       Any       		= self.dummy_model()
          snake_case__    :       Dict       		= self.dummy_sample_deter * scheduler.init_noise_sigma
          snake_case__    :       Union[str, Any]       		= sample.to(snake_case_		)

          for i, t in enumerate(scheduler.timesteps		):
                snake_case__    :       int       		= scheduler.scale_model_input(snake_case_					,	snake_case_		)

                snake_case__    :       List[Any]       		= model(snake_case_					,	snake_case_		)

                snake_case__    :       str       		= scheduler.step(snake_case_					,	snake_case_					,	snake_case_		)
                snake_case__    :       Optional[int]       		= output.prev_sample

          snake_case__    :       List[str]       		= torch.sum(torch.abs(snake_case_		)		)
          snake_case__    :       Union[str, Any]       		= torch.mean(torch.abs(snake_case_		)		)

          if torch_device in ["cpu", "mps"]:
                assert abs(result_sum.item() - 20.4125		) < 1E-2
                assert abs(result_mean.item() - 0.0266		) < 1E-3
          else:
                # CUDA
                assert abs(result_sum.item() - 20.4125		) < 1E-2
                assert abs(result_mean.item() - 0.0266		) < 1E-3




    def lowerCamelCase			(       self   :	Optional[int]		):
          if torch_device == "mps":
                return
          snake_case__    :       Any       		= self.scheduler_classes[0]
          snake_case__    :       str       		= self.get_scheduler_config()
          snake_case__    :       Dict       		= scheduler_class(**snake_case_		)

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

          snake_case__    :       List[str]       		= self.dummy_model()
          snake_case__    :       Optional[Any]       		= self.dummy_sample_deter.to(snake_case_		) * scheduler.init_noise_sigma

          for t in scheduler.timesteps:
                snake_case__    :       int       		= scheduler.scale_model_input(snake_case_					,	snake_case_		)

                snake_case__    :       Optional[int]       		= model(snake_case_					,	snake_case_		)

                snake_case__    :       int       		= scheduler.step(snake_case_					,	snake_case_					,	snake_case_		)
                snake_case__    :       str       		= output.prev_sample

          snake_case__    :       List[Any]       		= torch.sum(torch.abs(snake_case_		)		)
          snake_case__    :       Optional[int]       		= torch.mean(torch.abs(snake_case_		)		)

          if str(snake_case_		).startswith("""cpu"""		):
                # The following sum varies between 148 and 156 on mps. Why?
                assert abs(result_sum.item() - 20.4125		) < 1E-2
                assert abs(result_mean.item() - 0.0266		) < 1E-3
          else:
                # CUDA
                assert abs(result_sum.item() - 20.4125		) < 1E-2
                assert abs(result_mean.item() - 0.0266		) < 1E-3





'''simple docstring'''





import unittest

import torch
from torch import nn

from diffusers.models.activations import get_activation







class        UpperCAmelCase_			(			unittest.TestCase      ):
    """simple docstring"""







    def lowerCamelCase			(       self   :	Union[str, Any]		):
          snake_case__    :       int       		= get_activation("""swish"""		)

          self.assertIsInstance(snake_case_					,	nn.SiLU		)

          self.assertEqual(act(torch.tensor(-100					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertNotEqual(act(torch.tensor(-1					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertEqual(act(torch.tensor(0					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertEqual(act(torch.tensor(20					,	dtype=torch.floataa		)		).item()					,	20		)




    def lowerCamelCase			(       self   :	Optional[Any]		):
          snake_case__    :       int       		= get_activation("""silu"""		)

          self.assertIsInstance(snake_case_					,	nn.SiLU		)

          self.assertEqual(act(torch.tensor(-100					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertNotEqual(act(torch.tensor(-1					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertEqual(act(torch.tensor(0					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertEqual(act(torch.tensor(20					,	dtype=torch.floataa		)		).item()					,	20		)




    def lowerCamelCase			(       self   :	Dict		):
          snake_case__    :       Union[str, Any]       		= get_activation("""mish"""		)

          self.assertIsInstance(snake_case_					,	nn.Mish		)

          self.assertEqual(act(torch.tensor(-200					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertNotEqual(act(torch.tensor(-1					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertEqual(act(torch.tensor(0					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertEqual(act(torch.tensor(20					,	dtype=torch.floataa		)		).item()					,	20		)




    def lowerCamelCase			(       self   :	Union[str, Any]		):
          snake_case__    :       int       		= get_activation("""gelu"""		)

          self.assertIsInstance(snake_case_					,	nn.GELU		)

          self.assertEqual(act(torch.tensor(-100					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertNotEqual(act(torch.tensor(-1					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertEqual(act(torch.tensor(0					,	dtype=torch.floataa		)		).item()					,	0		)
          self.assertEqual(act(torch.tensor(20					,	dtype=torch.floataa		)		).item()					,	20		)
