code: import torch
  
  from diffusers import HeunDiscreteScheduler
  from diffusers.utils import torch_device
  
  from .test_schedulers import SchedulerCommonTest
  
  
  class _lowercase					(       _UpperCAmelCase						):
  
  
  
  
  
  							"""simple docstring"""
  
  
  
  
  							lowerCAmelCase__       =					(HeunDiscreteScheduler,)
  							lowerCAmelCase__       =					10
  
  
  
  							def     _UpperCAmelCase       (    self     ,   **UpperCAmelCase  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														_lowercase				=	{
  														    """num_train_timesteps""": 1100,
  														    """beta_start""": 0.0_001,
  														    """beta_end""": 0.02,
  														    """beta_schedule""": """linear""",
  														}
  
  														config.update(**UpperCAmelCase  )
  														return config
  
  
  
  							def     _UpperCAmelCase       (    self  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														for timesteps in [10, 50, 100, 1000]:
  																					self.check_over_configs(num_train_timesteps=UpperCAmelCase  )
  
  
  
  							def     _UpperCAmelCase       (    self  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														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=UpperCAmelCase     ,   beta_end=UpperCAmelCase  )
  
  
  
  							def     _UpperCAmelCase       (    self  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														for schedule in ["linear", "scaled_linear", "exp"]:
  																					self.check_over_configs(beta_schedule=UpperCAmelCase  )
  
  
  
  							def     _UpperCAmelCase       (    self  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														for clip_sample_range in [1.0, 2.0, 3.0]:
  																					self.check_over_configs(clip_sample_range=UpperCAmelCase     ,   clip_sample=UpperCAmelCase  )
  
  
  
  							def     _UpperCAmelCase       (    self  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														for prediction_type in ["epsilon", "v_prediction", "sample"]:
  																					self.check_over_configs(prediction_type=UpperCAmelCase  )
  
  
  
  							def     _UpperCAmelCase       (    self  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														_lowercase				=	self.scheduler_classes[0]
  														_lowercase				=	self.get_scheduler_config()
  														_lowercase				=	scheduler_class(**UpperCAmelCase  )
  
  														scheduler.set_timesteps(self.num_inference_steps  )
  
  														_lowercase				=	self.dummy_model()
  														_lowercase				=	self.dummy_sample_deter * scheduler.init_noise_sigma
  														_lowercase				=	sample.to(UpperCAmelCase  )
  
  														for i, t in enumerate(scheduler.timesteps  ):
  																					_lowercase				=	scheduler.scale_model_input(UpperCAmelCase     ,   UpperCAmelCase  )
  
  																					_lowercase				=	model(UpperCAmelCase     ,   UpperCAmelCase  )
  
  																					_lowercase				=	scheduler.step(UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  )
  																					_lowercase				=	output.prev_sample
  
  														_lowercase				=	torch.sum(torch.abs(UpperCAmelCase  )  )
  														_lowercase				=	torch.mean(torch.abs(UpperCAmelCase  )  )
  
  														if torch_device in ["cpu", "mps"]:
  																					assert abs(result_sum.item() - 0.1_233  ) < 1e-2
  																					assert abs(result_mean.item() - 0.0_002  ) < 1e-3
  														else:
  																					# CUDA
  																					assert abs(result_sum.item() - 0.1_233  ) < 1e-2
  																					assert abs(result_mean.item() - 0.0_002  ) < 1e-3
  
  
  
  							def     _UpperCAmelCase       (    self  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														_lowercase				=	self.scheduler_classes[0]
  														_lowercase				=	self.get_scheduler_config(prediction_type="""v_prediction"""  )
  														_lowercase				=	scheduler_class(**UpperCAmelCase  )
  
  														scheduler.set_timesteps(self.num_inference_steps  )
  
  														_lowercase				=	self.dummy_model()
  														_lowercase				=	self.dummy_sample_deter * scheduler.init_noise_sigma
  														_lowercase				=	sample.to(UpperCAmelCase  )
  
  														for i, t in enumerate(scheduler.timesteps  ):
  																					_lowercase				=	scheduler.scale_model_input(UpperCAmelCase     ,   UpperCAmelCase  )
  
  																					_lowercase				=	model(UpperCAmelCase     ,   UpperCAmelCase  )
  
  																					_lowercase				=	scheduler.step(UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  )
  																					_lowercase				=	output.prev_sample
  
  														_lowercase				=	torch.sum(torch.abs(UpperCAmelCase  )  )
  														_lowercase				=	torch.mean(torch.abs(UpperCAmelCase  )  )
  
  														if torch_device in ["cpu", "mps"]:
  																					assert abs(result_sum.item() - 4.6934e-07  ) < 1e-2
  																					assert abs(result_mean.item() - 6.1112e-10  ) < 1e-3
  														else:
  																					# CUDA
  																					assert abs(result_sum.item() - 4.693_4286_5017_0972e-07  ) < 1e-2
  																					assert abs(result_mean.item() - 0.0_002  ) < 1e-3
  
  
  
  							def     _UpperCAmelCase       (    self  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														_lowercase				=	self.scheduler_classes[0]
  														_lowercase				=	self.get_scheduler_config()
  														_lowercase				=	scheduler_class(**UpperCAmelCase  )
  
  														scheduler.set_timesteps(self.num_inference_steps     ,   device=UpperCAmelCase  )
  
  														_lowercase				=	self.dummy_model()
  														_lowercase				=	self.dummy_sample_deter.to(UpperCAmelCase  ) * scheduler.init_noise_sigma
  
  														for t in scheduler.timesteps:
  																					_lowercase				=	scheduler.scale_model_input(UpperCAmelCase     ,   UpperCAmelCase  )
  
  																					_lowercase				=	model(UpperCAmelCase     ,   UpperCAmelCase  )
  
  																					_lowercase				=	scheduler.step(UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  )
  																					_lowercase				=	output.prev_sample
  
  														_lowercase				=	torch.sum(torch.abs(UpperCAmelCase  )  )
  														_lowercase				=	torch.mean(torch.abs(UpperCAmelCase  )  )
  
  														if str(UpperCAmelCase  ).startswith("""cpu"""  ):
  																					# The following sum varies between 148 and 156 on mps. Why?
  																					assert abs(result_sum.item() - 0.1_233  ) < 1e-2
  																					assert abs(result_mean.item() - 0.0_002  ) < 1e-3
  														elif str(UpperCAmelCase  ).startswith("""mps"""  ):
  																					# Larger tolerance on mps
  																					assert abs(result_mean.item() - 0.0_002  ) < 1e-2
  														else:
  																					# CUDA
  																					assert abs(result_sum.item() - 0.1_233  ) < 1e-2
  																					assert abs(result_mean.item() - 0.0_002  ) < 1e-3
  
  
  
  							def     _UpperCAmelCase       (    self  ):
  
  
  
  
  
  
  
  														'''simple docstring'''
  
  
  
  
  
  
  														_lowercase				=	self.scheduler_classes[0]
  														_lowercase				=	self.get_scheduler_config()
  														_lowercase				=	scheduler_class(**UpperCAmelCase     ,   use_karras_sigmas=UpperCAmelCase  )
  
  														scheduler.set_timesteps(self.num_inference_steps     ,   device=UpperCAmelCase  )
  
  														_lowercase				=	self.dummy_model()
  														_lowercase				=	self.dummy_sample_deter.to(UpperCAmelCase  ) * scheduler.init_noise_sigma
  														_lowercase				=	sample.to(UpperCAmelCase  )
  
  														for t in scheduler.timesteps:
  																					_lowercase				=	scheduler.scale_model_input(UpperCAmelCase     ,   UpperCAmelCase  )
  
  																					_lowercase				=	model(UpperCAmelCase     ,   UpperCAmelCase  )
  
  																					_lowercase				=	scheduler.step(UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  )
  																					_lowercase				=	output.prev_sample
  
  														_lowercase				=	torch.sum(torch.abs(UpperCAmelCase  )  )
  														_lowercase				=	torch.mean(torch.abs(UpperCAmelCase  )  )
  
  														assert abs(result_sum.item() - 0.00_015  ) < 1e-2
  														assert abs(result_mean.item() - 1.9869_5545_3503_4695e-07  ) < 1e-2
  
  
  
style_context: from typing import TYPE_CHECKING
  
  from ...utils import (
      OptionalDependencyNotAvailable,
      _LazyModule,
      is_flax_available,
      is_tf_available,
      is_tokenizers_available,
      is_torch_available,
  )
  
  
  A_:      int        =			{
      'configuration_bart': ['BART_PRETRAINED_CONFIG_ARCHIVE_MAP', 'BartConfig', 'BartOnnxConfig'],
      'tokenization_bart': ['BartTokenizer'],
  }
  
  try:
  				if not is_tokenizers_available():
  								raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
  				pass
  else:
  				A_:      List[str]        =			['BartTokenizerFast']
  
  try:
  				if not is_torch_available():
  								raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
  				pass
  else:
  				A_:      int        =			[
  				    'BART_PRETRAINED_MODEL_ARCHIVE_LIST',
  				    'BartForCausalLM',
  				    'BartForConditionalGeneration',
  				    'BartForQuestionAnswering',
  				    'BartForSequenceClassification',
  				    'BartModel',
  				    'BartPretrainedModel',
  				    'PretrainedBartModel',
  				]
  
  try:
  				if not is_tf_available():
  								raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
  				pass
  else:
  				A_:      Any        =			[
  				    'TFBartForConditionalGeneration',
  				    'TFBartForSequenceClassification',
  				    'TFBartModel',
  				    'TFBartPretrainedModel',
  				]
  
  try:
  				if not is_flax_available():
  								raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
  				pass
  else:
  				A_:      str        =			[
  				    'FlaxBartDecoderPreTrainedModel',
  				    'FlaxBartForCausalLM',
  				    'FlaxBartForConditionalGeneration',
  				    'FlaxBartForQuestionAnswering',
  				    'FlaxBartForSequenceClassification',
  				    'FlaxBartModel',
  				    'FlaxBartPreTrainedModel',
  				]
  
  if TYPE_CHECKING:
  				from .configuration_bart import BART_PRETRAINED_CONFIG_ARCHIVE_MAP, BartConfig, BartOnnxConfig
  				from .tokenization_bart import BartTokenizer
  
  				try:
  								if not is_tokenizers_available():
  												raise OptionalDependencyNotAvailable()
  				except OptionalDependencyNotAvailable:
  								pass
  				else:
  								from .tokenization_bart_fast import BartTokenizerFast
  
  				try:
  								if not is_torch_available():
  												raise OptionalDependencyNotAvailable()
  				except OptionalDependencyNotAvailable:
  								pass
  				else:
  								from .modeling_bart import (
  								    BART_PRETRAINED_MODEL_ARCHIVE_LIST,
  								    BartForCausalLM,
  								    BartForConditionalGeneration,
  								    BartForQuestionAnswering,
  								    BartForSequenceClassification,
  								    BartModel,
  								    BartPretrainedModel,
  								    PretrainedBartModel,
  								)
  
  				try:
  								if not is_tf_available():
  												raise OptionalDependencyNotAvailable()
  				except OptionalDependencyNotAvailable:
  								pass
  				else:
  								from .modeling_tf_bart import (
  								    TFBartForConditionalGeneration,
  								    TFBartForSequenceClassification,
  								    TFBartModel,
  								    TFBartPretrainedModel,
  								)
  
  				try:
  								if not is_flax_available():
  												raise OptionalDependencyNotAvailable()
  				except OptionalDependencyNotAvailable:
  								pass
  				else:
  								from .modeling_flax_bart import (
  								    FlaxBartDecoderPreTrainedModel,
  								    FlaxBartForCausalLM,
  								    FlaxBartForConditionalGeneration,
  								    FlaxBartForQuestionAnswering,
  								    FlaxBartForSequenceClassification,
  								    FlaxBartModel,
  								    FlaxBartPreTrainedModel,
  								)
  
  else:
  				import sys
  
  				A_:      Optional[Any]        =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
  
  
  