





import torch

from diffusers import DPMSolverSDEScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import require_torchsde

from .test_schedulers import SchedulerCommonTest
@require_torchsde
class      UpperCamelCase_						(					_lowerCamelCase   ):
     lowerCAmelCase_							=		(DPMSolverSDEScheduler,)
     lowerCAmelCase_							=		10




     def 		lowerCAmelCase     (						self			,  **lowerCAmelCase_					)						->							Union[str, Any]:
          _snake_case													=			{
              'num_train_timesteps': 1100,
              'beta_start': 0.00_01,
              'beta_end': 0.02,
              'beta_schedule': 'linear',
              'noise_sampler_seed': 0,
          }

          config.update(**lowerCAmelCase_					)
          return config




     def 		lowerCAmelCase     (						self					)						->							Optional[Any]:
          for timesteps in [10, 50, 100, 1000]:
               self.check_over_configs(num_train_timesteps=lowerCAmelCase_					)




     def 		lowerCAmelCase     (						self					)						->							Optional[int]:
          for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]			,  [0.00_02, 0.0_02, 0.02]					):
               self.check_over_configs(beta_start=lowerCAmelCase_			,  beta_end=lowerCAmelCase_					)




     def 		lowerCAmelCase     (						self					)						->							int:
          for schedule in ["linear", "scaled_linear"]:
               self.check_over_configs(beta_schedule=lowerCAmelCase_					)




     def 		lowerCAmelCase     (						self					)						->							List[str]:
          for prediction_type in ["epsilon", "v_prediction"]:
               self.check_over_configs(prediction_type=lowerCAmelCase_					)




     def 		lowerCAmelCase     (						self					)						->							List[str]:
          _snake_case													=			self.scheduler_classes[0]
          _snake_case													=			self.get_scheduler_config()
          _snake_case													=			scheduler_class(**lowerCAmelCase_					)

          scheduler.set_timesteps(self.num_inference_steps					)

          _snake_case													=			self.dummy_model()
          _snake_case													=			self.dummy_sample_deter * scheduler.init_noise_sigma
          _snake_case													=			sample.to(lowerCAmelCase_					)

          for i, t in enumerate(scheduler.timesteps					):
               _snake_case													=			scheduler.scale_model_input(lowerCAmelCase_			,  lowerCAmelCase_					)

               _snake_case													=			model(lowerCAmelCase_			,  lowerCAmelCase_					)

               _snake_case													=			scheduler.step(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)
               _snake_case													=			output.prev_sample

          _snake_case													=			torch.sum(torch.abs(lowerCAmelCase_					)					)
          _snake_case													=			torch.mean(torch.abs(lowerCAmelCase_					)					)

          if torch_device in ["mps"]:
               assert abs(result_sum.item() - 167.47_8210_4492_1875					) < 1E-2
               assert abs(result_mean.item() - 0.21_78_70_59_64_56_52_77					) < 1E-3
          elif torch_device in ["cuda"]:
               assert abs(result_sum.item() - 171.59_3521_1181_6406					) < 1E-2
               assert abs(result_mean.item() - 0.2_23_42_90_68_92_29_96_52					) < 1E-3
          else:
               assert abs(result_sum.item() - 162.52_3834_2285_1562					) < 1E-2
               assert abs(result_mean.item() - 0.2_11_61_95_70_85_13_26					) < 1E-3




     def 		lowerCAmelCase     (						self					)						->							Optional[Any]:
          _snake_case													=			self.scheduler_classes[0]
          _snake_case													=			self.get_scheduler_config(prediction_type='v_prediction'					)
          _snake_case													=			scheduler_class(**lowerCAmelCase_					)

          scheduler.set_timesteps(self.num_inference_steps					)

          _snake_case													=			self.dummy_model()
          _snake_case													=			self.dummy_sample_deter * scheduler.init_noise_sigma
          _snake_case													=			sample.to(lowerCAmelCase_					)

          for i, t in enumerate(scheduler.timesteps					):
               _snake_case													=			scheduler.scale_model_input(lowerCAmelCase_			,  lowerCAmelCase_					)

               _snake_case													=			model(lowerCAmelCase_			,  lowerCAmelCase_					)

               _snake_case													=			scheduler.step(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)
               _snake_case													=			output.prev_sample

          _snake_case													=			torch.sum(torch.abs(lowerCAmelCase_					)					)
          _snake_case													=			torch.mean(torch.abs(lowerCAmelCase_					)					)

          if torch_device in ["mps"]:
               assert abs(result_sum.item() - 124.77_1492_0043_9453					) < 1E-2
               assert abs(result_mean.item() - 0.1_62_26_28_90_14_81_62_84					) < 1E-3
          elif torch_device in ["cuda"]:
               assert abs(result_sum.item() - 128.1_6633_6059_5703					) < 1E-2
               assert abs(result_mean.item() - 0.1_66_88_32_60_01_16_72_97					) < 1E-3
          else:
               assert abs(result_sum.item() - 119.8_4875_4882_8125					) < 1E-2
               assert abs(result_mean.item() - 0.15_60_53_06_62_53_66_21					) < 1E-3




     def 		lowerCAmelCase     (						self					)						->							Optional[int]:
          _snake_case													=			self.scheduler_classes[0]
          _snake_case													=			self.get_scheduler_config()
          _snake_case													=			scheduler_class(**lowerCAmelCase_					)

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

          _snake_case													=			self.dummy_model()
          _snake_case													=			self.dummy_sample_deter.to(lowerCAmelCase_					) * scheduler.init_noise_sigma

          for t in scheduler.timesteps:
               _snake_case													=			scheduler.scale_model_input(lowerCAmelCase_			,  lowerCAmelCase_					)

               _snake_case													=			model(lowerCAmelCase_			,  lowerCAmelCase_					)

               _snake_case													=			scheduler.step(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)
               _snake_case													=			output.prev_sample

          _snake_case													=			torch.sum(torch.abs(lowerCAmelCase_					)					)
          _snake_case													=			torch.mean(torch.abs(lowerCAmelCase_					)					)

          if torch_device in ["mps"]:
               assert abs(result_sum.item() - 167.46_9573_9746_0938					) < 1E-2
               assert abs(result_mean.item() - 0.2_18_05_93_46_07_98_26_35					) < 1E-3
          elif torch_device in ["cuda"]:
               assert abs(result_sum.item() - 171.59_3536_3769_5312					) < 1E-2
               assert abs(result_mean.item() - 0.2_23_42_90_83_82_41_57_71					) < 1E-3
          else:
               assert abs(result_sum.item() - 162.52_3834_2285_1562					) < 1E-2
               assert abs(result_mean.item() - 0.2_11_61_95_70_85_13_26					) < 1E-3






     def 		lowerCAmelCase     (						self					)						->							Optional[int]:
          _snake_case													=			self.scheduler_classes[0]
          _snake_case													=			self.get_scheduler_config()
          _snake_case													=			scheduler_class(**lowerCAmelCase_			,  use_karras_sigmas=lowerCAmelCase_					)

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

          _snake_case													=			self.dummy_model()
          _snake_case													=			self.dummy_sample_deter.to(lowerCAmelCase_					) * scheduler.init_noise_sigma
          _snake_case													=			sample.to(lowerCAmelCase_					)

          for t in scheduler.timesteps:
               _snake_case													=			scheduler.scale_model_input(lowerCAmelCase_			,  lowerCAmelCase_					)

               _snake_case													=			model(lowerCAmelCase_			,  lowerCAmelCase_					)

               _snake_case													=			scheduler.step(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_					)
               _snake_case													=			output.prev_sample

          _snake_case													=			torch.sum(torch.abs(lowerCAmelCase_					)					)
          _snake_case													=			torch.mean(torch.abs(lowerCAmelCase_					)					)

          if torch_device in ["mps"]:
               assert abs(result_sum.item() - 176.66_9741_3574_2188					) < 1E-2
               assert abs(result_mean.item() - 0.2_30_03_87_27_30_98_18_11					) < 1E-2
          elif torch_device in ["cuda"]:
               assert abs(result_sum.item() - 177.63_6535_6445_3125					) < 1E-2
               assert abs(result_mean.item() - 0.2_30_03_87_27_30_98_18_11					) < 1E-2
          else:
               assert abs(result_sum.item() - 170.3_1352_2338_8672					) < 1E-2
               assert abs(result_mean.item() - 0.2_30_03_87_27_30_98_18_11					) < 1E-2










import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch
@require_torch
class      UpperCamelCase_						:




					@property
					def 		lowerCAmelCase     (						self					)						->							int:
										return self.get_dummy_input()




					@property
					def 		lowerCAmelCase     (						self					)						->							Optional[Any]:
										if self.block_type == "down":
															return (4, 32, 16, 16)
										elif self.block_type == "mid":
															return (4, 32, 32, 32)
										elif self.block_type == "up":
															return (4, 32, 64, 64)

										raise ValueError(F'''\'{self.block_type}\' is not a supported block_type. Set it to \'up\', \'mid\', or \'down\'.'''					)




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_=True			,  lowerCAmelCase_=False			,  lowerCAmelCase_=False			,  lowerCAmelCase_=False			,  )						->							List[str]:
										_snake_case													=			4
										_snake_case													=			32
										_snake_case													=			(32, 32)

										_snake_case													=			torch.manual_seed(0					)
										_snake_case													=			torch.device(lowerCAmelCase_					)
										_snake_case													=			(batch_size, num_channels) + sizes
										_snake_case													=			randn_tensor(lowerCAmelCase_			,  generator=lowerCAmelCase_			,  device=lowerCAmelCase_					)
										_snake_case													=			{'hidden_states': hidden_states}

										if include_temb:
															_snake_case													=			128
															_snake_case													=			randn_tensor((batch_size, temb_channels)			,  generator=lowerCAmelCase_			,  device=lowerCAmelCase_					)

										if include_res_hidden_states_tuple:
															_snake_case													=			torch.manual_seed(1					)
															_snake_case													=			(randn_tensor(lowerCAmelCase_			,  generator=lowerCAmelCase_			,  device=lowerCAmelCase_					),)

										if include_encoder_hidden_states:
															_snake_case													=			floats_tensor((batch_size, 32, 32)					).to(lowerCAmelCase_					)

										if include_skip_sample:
															_snake_case													=			randn_tensor(((batch_size, 3) + sizes)			,  generator=lowerCAmelCase_			,  device=lowerCAmelCase_					)

										return dummy_input




					def 		lowerCAmelCase     (						self					)						->							Tuple:
										_snake_case													=			{
										    'in_channels': 32,
										    'out_channels': 32,
										    'temb_channels': 128,
										}
										if self.block_type == "up":
															_snake_case													=			32

										if self.block_type == "mid":
															init_dict.pop('out_channels'					)

										_snake_case													=			self.dummy_input
										return init_dict, inputs_dict




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_					)						->							Optional[int]:
										_snake_case							,       _snake_case													=			self.prepare_init_args_and_inputs_for_common()
										_snake_case													=			self.block_class(**lowerCAmelCase_					)
										unet_block.to(lowerCAmelCase_					)
										unet_block.eval()

										with torch.no_grad():
															_snake_case													=			unet_block(**lowerCAmelCase_					)

										if isinstance(lowerCAmelCase_			,  lowerCAmelCase_					):
															_snake_case													=			output[0]

										self.assertEqual(output.shape			,  self.output_shape					)

										_snake_case													=			output[0, -1, -3:, -3:]
										_snake_case													=			torch.tensor(lowerCAmelCase_					).to(lowerCAmelCase_					)
										assert torch_all_close(output_slice.flatten()			,  lowerCAmelCase_			,  atol=5E-3					)






					@unittest.skipIf(torch_device == 'mps'			,  'Training is not supported in mps'					)
					def 		lowerCAmelCase     (						self					)						->							Tuple:
										_snake_case							,       _snake_case													=			self.prepare_init_args_and_inputs_for_common()
										_snake_case													=			self.block_class(**lowerCAmelCase_					)
										model.to(lowerCAmelCase_					)
										model.train()
										_snake_case													=			model(**lowerCAmelCase_					)

										if isinstance(lowerCAmelCase_			,  lowerCAmelCase_					):
															_snake_case													=			output[0]

										_snake_case													=			torch.device(lowerCAmelCase_					)
										_snake_case													=			randn_tensor(output.shape			,  device=lowerCAmelCase_					)
										_snake_case													=			torch.nn.functional.mse_loss(lowerCAmelCase_			,  lowerCAmelCase_					)
										loss.backward()



