
import flax.linen as nn
import jax
import jax.numpy as jnp


class 			lowercase_		(      nn.Module       ):
       __lowerCamelCase				       =						42
       __lowerCamelCase				       =						jnp.floataa


       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=nn.Conv(
                  self.out_channels						,       kernel_size=(3, 3)						,       strides=(1, 1)						,       padding=((1, 1), (1, 1))						,       dtype=self.dtype						,       )





       def __call__(    self						,       __A				)						->    int:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Optional[int]		=hidden_states.shape
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=jax.image.resize(
                  __A						,       shape=(batch, height * 2, width * 2, channels)						,       method='''nearest'''						,       )
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.conv(__A				)
              return hidden_states






class 			lowercase_		(      nn.Module       ):
       __lowerCamelCase				       =						42
       __lowerCamelCase				       =						jnp.floataa


       def 	_snake_case				(    self				)						->    List[str]:
              SCREAMING_SNAKE_CASE_							:				str		=nn.Conv(
                  self.out_channels						,       kernel_size=(3, 3)						,       strides=(2, 2)						,       padding=((1, 1), (1, 1))						,       dtype=self.dtype						,       )





       def __call__(    self						,       __A				)						->    str:
              # pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
              # hidden_states = jnp.pad(hidden_states, pad_width=pad)
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=self.conv(__A				)
              return hidden_states






class 			lowercase_		(      nn.Module       ):
       __lowerCamelCase				       =						42
       __lowerCamelCase				       =						None
       __lowerCamelCase				       =						0.0
       __lowerCamelCase				       =						None
       __lowerCamelCase				       =						jnp.floataa


       def 	_snake_case				(    self				)						->    Optional[int]:
              SCREAMING_SNAKE_CASE_							:				Any		=self.in_channels if self.out_channels is None else self.out_channels

              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=nn.GroupNorm(num_groups=32						,       epsilon=1e-5				)
              SCREAMING_SNAKE_CASE_							:				int		=nn.Conv(
                  __A						,       kernel_size=(3, 3)						,       strides=(1, 1)						,       padding=((1, 1), (1, 1))						,       dtype=self.dtype						,       )

              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=nn.Dense(__A						,       dtype=self.dtype				)

              SCREAMING_SNAKE_CASE_							:				Any		=nn.GroupNorm(num_groups=32						,       epsilon=1e-5				)
              SCREAMING_SNAKE_CASE_							:				str		=nn.Dropout(self.dropout_prob				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=nn.Conv(
                  __A						,       kernel_size=(3, 3)						,       strides=(1, 1)						,       padding=((1, 1), (1, 1))						,       dtype=self.dtype						,       )

              SCREAMING_SNAKE_CASE_							:				Any		=self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=None
              if use_nin_shortcut:
                     SCREAMING_SNAKE_CASE_							:				Any		=nn.Conv(
                         __A						,       kernel_size=(1, 1)						,       strides=(1, 1)						,       padding='''VALID'''						,       dtype=self.dtype						,       )





       def __call__(    self						,       __A						,       __A						,       __A=True				)						->    Optional[int]:
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=hidden_states
              SCREAMING_SNAKE_CASE_							:				int		=self.norma(__A				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=nn.swish(__A				)
              SCREAMING_SNAKE_CASE_							:				Dict		=self.conva(__A				)

              SCREAMING_SNAKE_CASE_							:				List[str]		=self.time_emb_proj(nn.swish(__A				)				)
              SCREAMING_SNAKE_CASE_							:				int		=jnp.expand_dims(jnp.expand_dims(__A						,       1				)						,       1				)
              SCREAMING_SNAKE_CASE_							:				Tuple		=hidden_states + temb

              SCREAMING_SNAKE_CASE_							:				Dict		=self.norma(__A				)
              SCREAMING_SNAKE_CASE_							:				int		=nn.swish(__A				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=self.dropout(__A						,       __A				)
              SCREAMING_SNAKE_CASE_							:				Dict		=self.conva(__A				)

              if self.conv_shortcut is not None:
                     SCREAMING_SNAKE_CASE_							:				List[str]		=self.conv_shortcut(__A				)

              return hidden_states + residual


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 			lowercase_		(      A       ):
       __lowerCamelCase				       =						(DPMSolverSDEScheduler,)
       __lowerCamelCase				       =						1_0


       def 	_snake_case				(    self						,       **__A				)						->    Union[str, Any]:
              SCREAMING_SNAKE_CASE_							:				Optional[int]		={
                  '''num_train_timesteps''': 1_100,
                  '''beta_start''': 0.0_001,
                  '''beta_end''': 0.02,
                  '''beta_schedule''': '''linear''',
                  '''noise_sampler_seed''': 0,
              }

              config.update(**__A				)
              return config


       def 	_snake_case				(    self				)						->    Any:
              for timesteps in [10, 50, 100, 1_000]:
                     self.check_over_configs(num_train_timesteps=__A				)


       def 	_snake_case				(    self				)						->    Union[str, 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=__A						,       beta_end=__A				)


       def 	_snake_case				(    self				)						->    Tuple:
              for schedule in ["linear", "scaled_linear"]:
                     self.check_over_configs(beta_schedule=__A				)


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


       def 	_snake_case				(    self				)						->    int:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=self.scheduler_classes[0]
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.get_scheduler_config()
              SCREAMING_SNAKE_CASE_							:				List[str]		=scheduler_class(**__A				)

              scheduler.set_timesteps(self.num_inference_steps				)

              SCREAMING_SNAKE_CASE_							:				int		=self.dummy_model()
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.dummy_sample_deter * scheduler.init_noise_sigma
              SCREAMING_SNAKE_CASE_							:				int		=sample.to(__A				)

              for i, t in enumerate(scheduler.timesteps				):
                     SCREAMING_SNAKE_CASE_							:				List[Any]		=scheduler.scale_model_input(__A						,       __A				)

                     SCREAMING_SNAKE_CASE_							:				Dict		=model(__A						,       __A				)

                     SCREAMING_SNAKE_CASE_							:				Tuple		=scheduler.step(__A						,       __A						,       __A				)
                     SCREAMING_SNAKE_CASE_							:				str		=output.prev_sample

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=torch.sum(torch.abs(__A				)				)
              SCREAMING_SNAKE_CASE_							:				List[str]		=torch.mean(torch.abs(__A				)				)

              if torch_device in ["mps"]:
                     assert abs(result_sum.item() - 167.47_821_044_921_875				) < 1e-2
                     assert abs(result_mean.item() - 0.2_178_705_964_565_277				) < 1e-3
              elif torch_device in ["cuda"]:
                     assert abs(result_sum.item() - 171.59_352_111_816_406				) < 1e-2
                     assert abs(result_mean.item() - 0.22_342_906_892_299_652				) < 1e-3
              else:
                     assert abs(result_sum.item() - 162.52_383_422_851_562				) < 1e-2
                     assert abs(result_mean.item() - 0.211_619_570_851_326				) < 1e-3


       def 	_snake_case				(    self				)						->    Dict:
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.scheduler_classes[0]
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=self.get_scheduler_config(prediction_type='''v_prediction'''				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=scheduler_class(**__A				)

              scheduler.set_timesteps(self.num_inference_steps				)

              SCREAMING_SNAKE_CASE_							:				Tuple		=self.dummy_model()
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.dummy_sample_deter * scheduler.init_noise_sigma
              SCREAMING_SNAKE_CASE_							:				Tuple		=sample.to(__A				)

              for i, t in enumerate(scheduler.timesteps				):
                     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=scheduler.scale_model_input(__A						,       __A				)

                     SCREAMING_SNAKE_CASE_							:				str		=model(__A						,       __A				)

                     SCREAMING_SNAKE_CASE_							:				Any		=scheduler.step(__A						,       __A						,       __A				)
                     SCREAMING_SNAKE_CASE_							:				Tuple		=output.prev_sample

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=torch.sum(torch.abs(__A				)				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=torch.mean(torch.abs(__A				)				)

              if torch_device in ["mps"]:
                     assert abs(result_sum.item() - 124.77_149_200_439_453				) < 1e-2
                     assert abs(result_mean.item() - 0.16_226_289_014_816_284				) < 1e-3
              elif torch_device in ["cuda"]:
                     assert abs(result_sum.item() - 128.1_663_360_595_703				) < 1e-2
                     assert abs(result_mean.item() - 0.16_688_326_001_167_297				) < 1e-3
              else:
                     assert abs(result_sum.item() - 119.8_487_548_828_125				) < 1e-2
                     assert abs(result_mean.item() - 0.1_560_530_662_536_621				) < 1e-3


       def 	_snake_case				(    self				)						->    int:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=self.scheduler_classes[0]
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.get_scheduler_config()
              SCREAMING_SNAKE_CASE_							:				List[str]		=scheduler_class(**__A				)

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

              SCREAMING_SNAKE_CASE_							:				int		=self.dummy_model()
              SCREAMING_SNAKE_CASE_							:				int		=self.dummy_sample_deter.to(__A				) * scheduler.init_noise_sigma

              for t in scheduler.timesteps:
                     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=scheduler.scale_model_input(__A						,       __A				)

                     SCREAMING_SNAKE_CASE_							:				Dict		=model(__A						,       __A				)

                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=scheduler.step(__A						,       __A						,       __A				)
                     SCREAMING_SNAKE_CASE_							:				List[Any]		=output.prev_sample

              SCREAMING_SNAKE_CASE_							:				List[Any]		=torch.sum(torch.abs(__A				)				)
              SCREAMING_SNAKE_CASE_							:				int		=torch.mean(torch.abs(__A				)				)

              if torch_device in ["mps"]:
                     assert abs(result_sum.item() - 167.46_957_397_460_938				) < 1e-2
                     assert abs(result_mean.item() - 0.21_805_934_607_982_635				) < 1e-3
              elif torch_device in ["cuda"]:
                     assert abs(result_sum.item() - 171.59_353_637_695_312				) < 1e-2
                     assert abs(result_mean.item() - 0.22_342_908_382_415_771				) < 1e-3
              else:
                     assert abs(result_sum.item() - 162.52_383_422_851_562				) < 1e-2
                     assert abs(result_mean.item() - 0.211_619_570_851_326				) < 1e-3





       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=self.scheduler_classes[0]
              SCREAMING_SNAKE_CASE_							:				List[str]		=self.get_scheduler_config()
              SCREAMING_SNAKE_CASE_							:				Dict		=scheduler_class(**__A						,       use_karras_sigmas=__A				)

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

              SCREAMING_SNAKE_CASE_							:				Dict		=self.dummy_model()
              SCREAMING_SNAKE_CASE_							:				List[Any]		=self.dummy_sample_deter.to(__A				) * scheduler.init_noise_sigma
              SCREAMING_SNAKE_CASE_							:				str		=sample.to(__A				)

              for t in scheduler.timesteps:
                     SCREAMING_SNAKE_CASE_							:				Any		=scheduler.scale_model_input(__A						,       __A				)

                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=model(__A						,       __A				)

                     SCREAMING_SNAKE_CASE_							:				int		=scheduler.step(__A						,       __A						,       __A				)
                     SCREAMING_SNAKE_CASE_							:				List[str]		=output.prev_sample

              SCREAMING_SNAKE_CASE_							:				str		=torch.sum(torch.abs(__A				)				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=torch.mean(torch.abs(__A				)				)

              if torch_device in ["mps"]:
                     assert abs(result_sum.item() - 176.66_974_135_742_188				) < 1e-2
                     assert abs(result_mean.item() - 0.23_003_872_730_981_811				) < 1e-2
              elif torch_device in ["cuda"]:
                     assert abs(result_sum.item() - 177.63_653_564_453_125				) < 1e-2
                     assert abs(result_mean.item() - 0.23_003_872_730_981_811				) < 1e-2
              else:
                     assert abs(result_sum.item() - 170.3_135_223_388_672				) < 1e-2
                     assert abs(result_mean.item() - 0.23_003_872_730_981_811				) < 1e-2
