




"""simple docstring"""




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPanoramaPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
@skip_mps
class 						_A (  a__      ,       a__      ,       unittest.TestCase				):
  _UpperCamelCase       :      Dict			=      StableDiffusionPanoramaPipeline
  _UpperCamelCase       :      Any			=      TEXT_TO_IMAGE_PARAMS
  _UpperCamelCase       :      Tuple			=      TEXT_TO_IMAGE_BATCH_PARAMS
  _UpperCamelCase       :      Any			=      TEXT_TO_IMAGE_IMAGE_PARAMS
  _UpperCamelCase       :      Union[str, Any]			=      TEXT_TO_IMAGE_IMAGE_PARAMS






  def  __a      (  self		:     Tuple				)		->   int:






       """simple docstring"""







       torch.manual_seed(0				)
       lowercase :       int       =			UNetaDConditionModel(
           block_out_channels=(32, 64)				,						layers_per_block=1				,						sample_size=32				,						in_channels=4				,						out_channels=4				,						down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,						up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,						cross_attention_dim=32				,						)
       lowercase :       int       =			DDIMScheduler()
       torch.manual_seed(0				)
       lowercase :       str       =			AutoencoderKL(
           block_out_channels=[32, 64]				,						in_channels=3				,						out_channels=3				,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']				,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']				,						latent_channels=4				,						)
       torch.manual_seed(0				)
       lowercase :       List[str]       =			CLIPTextConfig(
           bos_token_id=0				,						eos_token_id=2				,						hidden_size=32				,						intermediate_size=37				,						layer_norm_eps=1E-05				,						num_attention_heads=4				,						num_hidden_layers=5				,						pad_token_id=1				,						vocab_size=1_000				,						)
       lowercase :       Optional[int]       =			CLIPTextModel(lowerCAmelCase__				)
       lowercase :       Optional[int]       =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

       lowercase :       Optional[int]       =			{
           "unet": unet,
           "scheduler": scheduler,
           "vae": vae,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "safety_checker": None,
           "feature_extractor": None,
       }
       return components






  def  __a      (  self		:     Tuple				,						_A		:     List[Any]				,						_A		:     List[str]=0				)		->   Optional[Any]:






       """simple docstring"""







       lowercase :       int       =			torch.manual_seed(lowerCAmelCase__				)
       lowercase :       str       =			{
           "prompt": "a photo of the dolomites",
           "generator": generator,
           # Setting height and width to None to prevent OOMs on CPU.
           "height": None,
           "width": None,
           "num_inference_steps": 1,
           "guidance_scale": 6.0,
           "output_type": "numpy",
       }
       return inputs






  def  __a      (  self		:     Any				)		->   int:






       """simple docstring"""







       lowercase :       str       =			"cpu"  # ensure determinism for the device-dependent torch.Generator
       lowercase :       str       =			self.get_dummy_components()
       lowercase :       List[Any]       =			StableDiffusionPanoramaPipeline(**lowerCAmelCase__				)
       lowercase :       List[Any]       =			sd_pipe.to(lowerCAmelCase__				)
       sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

       lowercase :       Tuple       =			self.get_dummy_inputs(lowerCAmelCase__				)
       lowercase :       int       =			sd_pipe(**lowerCAmelCase__				).images
       lowercase :       Dict       =			image[0, -3:, -3:, -1]
       assert image.shape == (1, 64, 64, 3)

       lowercase :       str       =			np.array([0.6_186, 0.5_374, 0.4_915, 0.4_135, 0.4_114, 0.4_563, 0.5_128, 0.4_977, 0.4_757]				)

       assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2






  def  __a      (  self		:     int				)		->   Optional[int]:






       """simple docstring"""







       super().test_inference_batch_consistent(batch_sizes=[1, 2]				)






  def  __a      (  self		:     List[str]				)		->   Optional[int]:






       """simple docstring"""







       super().test_inference_batch_single_identical(batch_size=2				,						expected_max_diff=3.25E-3				)






  def  __a      (  self		:     Union[str, Any]				)		->   List[Any]:






       """simple docstring"""







       lowercase :       str       =			"cpu"  # ensure determinism for the device-dependent torch.Generator
       lowercase :       Dict       =			self.get_dummy_components()
       lowercase :       str       =			StableDiffusionPanoramaPipeline(**lowerCAmelCase__				)
       lowercase :       Tuple       =			sd_pipe.to(lowerCAmelCase__				)
       sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

       lowercase :       Tuple       =			self.get_dummy_inputs(lowerCAmelCase__				)
       lowercase :       Optional[int]       =			"french fries"
       lowercase :       Any       =			sd_pipe(**lowerCAmelCase__				,						negative_prompt=lowerCAmelCase__				)
       lowercase :       str       =			output.images
       lowercase :       List[str]       =			image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)

       lowercase :       Dict       =			np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757]				)

       assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2






  def  __a      (  self		:     Optional[Any]				)		->   Any:






       """simple docstring"""







       lowercase :       List[str]       =			"cpu"  # ensure determinism for the device-dependent torch.Generator
       lowercase :       str       =			self.get_dummy_components()
       lowercase :       Optional[Any]       =			StableDiffusionPanoramaPipeline(**lowerCAmelCase__				)
       lowercase :       int       =			sd_pipe.to(lowerCAmelCase__				)
       sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

       lowercase :       Union[str, Any]       =			self.get_dummy_inputs(lowerCAmelCase__				)
       lowercase :       List[str]       =			sd_pipe(**lowerCAmelCase__				,						view_batch_size=2				)
       lowercase :       Any       =			output.images
       lowercase :       Optional[int]       =			image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)

       lowercase :       Dict       =			np.array([0.6_187, 0.5_375, 0.4_915, 0.4_136, 0.4_114, 0.4_563, 0.5_128, 0.4_976, 0.4_757]				)

       assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2






  def  __a      (  self		:     str				)		->   Union[str, Any]:






       """simple docstring"""







       lowercase :       Optional[Any]       =			"cpu"  # ensure determinism for the device-dependent torch.Generator
       lowercase :       Tuple       =			self.get_dummy_components()
       lowercase :       List[Any]       =			EulerAncestralDiscreteScheduler(
           beta_start=0.00_085				,						beta_end=0.012				,						beta_schedule='''scaled_linear'''				)
       lowercase :       List[str]       =			StableDiffusionPanoramaPipeline(**lowerCAmelCase__				)
       lowercase :       int       =			sd_pipe.to(lowerCAmelCase__				)
       sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

       lowercase :       Union[str, Any]       =			self.get_dummy_inputs(lowerCAmelCase__				)
       lowercase :       str       =			sd_pipe(**lowerCAmelCase__				).images
       lowercase :       int       =			image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)

       lowercase :       str       =			np.array([0.4_024, 0.6_510, 0.4_901, 0.5_378, 0.5_813, 0.5_622, 0.4_795, 0.4_467, 0.4_952]				)

       assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2






  def  __a      (  self		:     List[str]				)		->   Any:






       """simple docstring"""







       lowercase :       Any       =			"cpu"  # ensure determinism for the device-dependent torch.Generator
       lowercase :       int       =			self.get_dummy_components()
       lowercase :       List[Any]       =			PNDMScheduler(
           beta_start=0.00_085				,						beta_end=0.012				,						beta_schedule='''scaled_linear'''				,						skip_prk_steps=lowerCAmelCase__				)
       lowercase :       Union[str, Any]       =			StableDiffusionPanoramaPipeline(**lowerCAmelCase__				)
       lowercase :       Any       =			sd_pipe.to(lowerCAmelCase__				)
       sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

       lowercase :       int       =			self.get_dummy_inputs(lowerCAmelCase__				)
       lowercase :       Optional[Any]       =			sd_pipe(**lowerCAmelCase__				).images
       lowercase :       Tuple       =			image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)

       lowercase :       Tuple       =			np.array([0.6_391, 0.6_291, 0.4_861, 0.5_134, 0.5_552, 0.4_578, 0.5_032, 0.5_023, 0.4_539]				)

       assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2


@slow
@require_torch_gpu
class 						_A (  unittest.TestCase				):






  def  __a      (  self		:     Dict				)		->   str:






       """simple docstring"""







       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






  def  __a      (  self		:     Optional[int]				,						_A		:     List[Any]=0				)		->   int:






       """simple docstring"""







       lowercase :       List[Any]       =			torch.manual_seed(lowerCAmelCase__				)
       lowercase :       Dict       =			{
           "prompt": "a photo of the dolomites",
           "generator": generator,
           "num_inference_steps": 3,
           "guidance_scale": 7.5,
           "output_type": "numpy",
       }
       return inputs






  def  __a      (  self		:     Union[str, Any]				)		->   List[str]:






       """simple docstring"""







       lowercase :       Any       =			"stabilityai/stable-diffusion-2-base"
       lowercase :       List[Any]       =			DDIMScheduler.from_pretrained(lowerCAmelCase__				,						subfolder='''scheduler'''				)
       lowercase :       List[Any]       =			StableDiffusionPanoramaPipeline.from_pretrained(lowerCAmelCase__				,						scheduler=lowerCAmelCase__				,						safety_checker=lowerCAmelCase__				)
       pipe.to(lowerCAmelCase__				)
       pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
       pipe.enable_attention_slicing()

       lowercase :       List[str]       =			self.get_inputs()
       lowercase :       Optional[int]       =			pipe(**lowerCAmelCase__				).images
       lowercase :       Optional[Any]       =			image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 512, 2_048, 3)

       lowercase :       Any       =			np.array(
           [
               0.36_968_392,
               0.27_025_372,
               0.32_446_766,
               0.28_379_387,
               0.36_363_274,
               0.30_733_347,
               0.27_100_027,
               0.27_054_125,
               0.25_536_096,
           ]				)

       assert np.abs(expected_slice - image_slice				).max() < 1E-2






  def  __a      (  self		:     List[str]				)		->   List[Any]:






       """simple docstring"""







       lowercase :       List[Any]       =			StableDiffusionPanoramaPipeline.from_pretrained(
           '''stabilityai/stable-diffusion-2-base'''				,						safety_checker=lowerCAmelCase__				)
       lowercase :       Any       =			LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
       pipe.to(lowerCAmelCase__				)
       pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
       pipe.enable_attention_slicing()

       lowercase :       Optional[Any]       =			self.get_inputs()
       lowercase :       Dict       =			pipe(**lowerCAmelCase__				).images
       lowercase :       Union[str, Any]       =			image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 512, 2_048, 3)

       lowercase :       List[Any]       =			np.array(
           [
               [
                   0.0,
                   0.0,
                   0.0,
                   0.0,
                   0.0,
                   0.0,
                   0.0,
                   0.0,
                   0.0,
               ]
           ]				)

       assert np.abs(expected_slice - image_slice				).max() < 1E-3






  def  __a      (  self		:     List[Any]				)		->   Dict:






       """simple docstring"""







       lowercase :       Tuple       =			0

       def callback_fn(_A		:     List[str]				,						_A		:     List[Any]				,						_A		:     Tuple				) -> None:
            lowercase :       Union[str, Any]       =			True
            nonlocal number_of_steps
            number_of_steps += 1
            if step == 1:
                 lowercase :       Any       =			latents.detach().cpu().numpy()
                 assert latents.shape == (1, 4, 64, 256)
                 lowercase :       Optional[int]       =			latents[0, -3:, -3:, -1]

                 lowercase :       Union[str, Any]       =			np.array(
                     [
                         0.18_681_869,
                         0.33_907_816,
                         0.5_361_276,
                         0.14_432_865,
                         -0.02_856_611,
                         -0.73_941_123,
                         0.23_397_987,
                         0.47_322_682,
                         -0.37_823_164,
                     ]				)
                 assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5E-2
            elif step == 2:
                 lowercase :       Any       =			latents.detach().cpu().numpy()
                 assert latents.shape == (1, 4, 64, 256)
                 lowercase :       Optional[int]       =			latents[0, -3:, -3:, -1]

                 lowercase :       str       =			np.array(
                     [
                         0.18_539_645,
                         0.33_987_248,
                         0.5_378_559,
                         0.14_437_142,
                         -0.02_455_261,
                         -0.7_338_317,
                         0.23_990_755,
                         0.47_356_272,
                         -0.3_786_505,
                     ]				)

                 assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5E-2

       lowercase :       Union[str, Any]       =			False

       lowercase :       Optional[int]       =			"stabilityai/stable-diffusion-2-base"
       lowercase :       Dict       =			DDIMScheduler.from_pretrained(lowerCAmelCase__				,						subfolder='''scheduler'''				)
       lowercase :       Optional[int]       =			StableDiffusionPanoramaPipeline.from_pretrained(lowerCAmelCase__				,						scheduler=lowerCAmelCase__				,						safety_checker=lowerCAmelCase__				)
       lowercase :       Optional[int]       =			pipe.to(lowerCAmelCase__				)
       pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
       pipe.enable_attention_slicing()

       lowercase :       Union[str, Any]       =			self.get_inputs()
       pipe(**lowerCAmelCase__				,						callback=lowerCAmelCase__				,						callback_steps=1				)
       assert callback_fn.has_been_called
       assert number_of_steps == 3






  def  __a      (  self		:     Optional[Any]				)		->   int:






       """simple docstring"""







       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

       lowercase :       Tuple       =			"stabilityai/stable-diffusion-2-base"
       lowercase :       Optional[int]       =			DDIMScheduler.from_pretrained(lowerCAmelCase__				,						subfolder='''scheduler'''				)
       lowercase :       int       =			StableDiffusionPanoramaPipeline.from_pretrained(lowerCAmelCase__				,						scheduler=lowerCAmelCase__				,						safety_checker=lowerCAmelCase__				)
       lowercase :       Dict       =			pipe.to(lowerCAmelCase__				)
       pipe.set_progress_bar_config(disable=lowerCAmelCase__				)
       pipe.enable_attention_slicing(1				)
       pipe.enable_sequential_cpu_offload()

       lowercase :       str       =			self.get_inputs()
       lowercase :       Optional[int]       =			pipe(**lowerCAmelCase__				)

       lowercase :       Optional[int]       =			torch.cuda.max_memory_allocated()
       # make sure that less than 5.2 GB is allocated
       assert mem_bytes < 5.5 * 10**9





import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest
class 						_A (  _lowerCamelCase				):
  _UpperCamelCase       :      Tuple			=      (PNDMScheduler,)
  _UpperCamelCase       :      Optional[int]			=      (('''num_inference_steps''', 5_0),)






  def  __a      (  self		:     Union[str, Any]				,						**_A		:     Optional[Any]				)		->   Optional[int]:






       """simple docstring"""







       lowercase :       List[str]       =			{
           '''num_train_timesteps''': 1_000,
           '''beta_start''': 0.0_001,
           '''beta_end''': 0.02,
           '''beta_schedule''': '''linear''',
       }

       config.update(**_A				)
       return config






  def  __a      (  self		:     Optional[Any]				,						_A		:     Any=0				,						**_A		:     int				)		->   Dict:






       """simple docstring"""







       lowercase :       List[Any]       =			dict(self.forward_default_kwargs				)
       lowercase :       Union[str, Any]       =			kwargs.pop('''num_inference_steps'''				,						_A				)
       lowercase :       Optional[Any]       =			self.dummy_sample
       lowercase :       str       =			0.1 * sample
       lowercase :       Union[str, Any]       =			[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

       for scheduler_class in self.scheduler_classes:
            lowercase :       Dict       =			self.get_scheduler_config(**_A				)
            lowercase :       int       =			scheduler_class(**_A				)
            scheduler.set_timesteps(_A				)
            # copy over dummy past residuals
            lowercase :       List[str]       =			dummy_past_residuals[:]

            with tempfile.TemporaryDirectory() as tmpdirname:
                 scheduler.save_config(_A				)
                 lowercase :       Union[str, Any]       =			scheduler_class.from_pretrained(_A				)
                 new_scheduler.set_timesteps(_A				)
                 # copy over dummy past residuals
                 lowercase :       int       =			dummy_past_residuals[:]

            lowercase :       List[str]       =			scheduler.step_prk(_A				,						_A				,						_A				,						**_A				).prev_sample
            lowercase :       List[Any]       =			new_scheduler.step_prk(_A				,						_A				,						_A				,						**_A				).prev_sample

            assert torch.sum(torch.abs(output - new_output				)				) < 1E-5, "Scheduler outputs are not identical"

            lowercase :       Dict       =			scheduler.step_plms(_A				,						_A				,						_A				,						**_A				).prev_sample
            lowercase :       List[Any]       =			new_scheduler.step_plms(_A				,						_A				,						_A				,						**_A				).prev_sample

            assert torch.sum(torch.abs(output - new_output				)				) < 1E-5, "Scheduler outputs are not identical"






  def  __a      (  self		:     int				)		->   List[str]:






       """simple docstring"""







       pass






  def  __a      (  self		:     Optional[Any]				,						_A		:     Optional[int]=0				,						**_A		:     str				)		->   Optional[Any]:






       """simple docstring"""







       lowercase :       Optional[int]       =			dict(self.forward_default_kwargs				)
       lowercase :       Optional[int]       =			kwargs.pop('''num_inference_steps'''				,						_A				)
       lowercase :       Union[str, Any]       =			self.dummy_sample
       lowercase :       Any       =			0.1 * sample
       lowercase :       int       =			[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

       for scheduler_class in self.scheduler_classes:
            lowercase :       List[str]       =			self.get_scheduler_config()
            lowercase :       List[str]       =			scheduler_class(**_A				)
            scheduler.set_timesteps(_A				)

            # copy over dummy past residuals (must be after setting timesteps)
            lowercase :       Union[str, Any]       =			dummy_past_residuals[:]

            with tempfile.TemporaryDirectory() as tmpdirname:
                 scheduler.save_config(_A				)
                 lowercase :       int       =			scheduler_class.from_pretrained(_A				)
                 # copy over dummy past residuals
                 new_scheduler.set_timesteps(_A				)

                 # copy over dummy past residual (must be after setting timesteps)
                 lowercase :       Optional[Any]       =			dummy_past_residuals[:]

            lowercase :       int       =			scheduler.step_prk(_A				,						_A				,						_A				,						**_A				).prev_sample
            lowercase :       Tuple       =			new_scheduler.step_prk(_A				,						_A				,						_A				,						**_A				).prev_sample

            assert torch.sum(torch.abs(output - new_output				)				) < 1E-5, "Scheduler outputs are not identical"

            lowercase :       Any       =			scheduler.step_plms(_A				,						_A				,						_A				,						**_A				).prev_sample
            lowercase :       Union[str, Any]       =			new_scheduler.step_plms(_A				,						_A				,						_A				,						**_A				).prev_sample

            assert torch.sum(torch.abs(output - new_output				)				) < 1E-5, "Scheduler outputs are not identical"






  def  __a      (  self		:     Union[str, Any]				,						**_A		:     Optional[int]				)		->   Any:






       """simple docstring"""







       lowercase :       str       =			self.scheduler_classes[0]
       lowercase :       Optional[int]       =			self.get_scheduler_config(**_A				)
       lowercase :       int       =			scheduler_class(**_A				)

       lowercase :       Any       =			10
       lowercase :       Optional[int]       =			self.dummy_model()
       lowercase :       Tuple       =			self.dummy_sample_deter
       scheduler.set_timesteps(_A				)

       for i, t in enumerate(scheduler.prk_timesteps				):
            lowercase :       List[str]       =			model(_A				,						_A				)
            lowercase :       Any       =			scheduler.step_prk(_A				,						_A				,						_A				).prev_sample

       for i, t in enumerate(scheduler.plms_timesteps				):
            lowercase :       Dict       =			model(_A				,						_A				)
            lowercase :       List[Any]       =			scheduler.step_plms(_A				,						_A				,						_A				).prev_sample

       return sample






  def  __a      (  self		:     List[Any]				)		->   List[str]:






       """simple docstring"""







       lowercase :       int       =			dict(self.forward_default_kwargs				)

       lowercase :       Union[str, Any]       =			kwargs.pop('''num_inference_steps'''				,						_A				)

       for scheduler_class in self.scheduler_classes:
            lowercase :       Optional[Any]       =			self.get_scheduler_config()
            lowercase :       List[Any]       =			scheduler_class(**_A				)

            lowercase :       Dict       =			self.dummy_sample
            lowercase :       str       =			0.1 * sample

            if num_inference_steps is not None and hasattr(_A				,						'''set_timesteps'''				):
                 scheduler.set_timesteps(_A				)
            elif num_inference_steps is not None and not hasattr(_A				,						'''set_timesteps'''				):
                 lowercase :       List[str]       =			num_inference_steps

            # copy over dummy past residuals (must be done after set_timesteps)
            lowercase :       Any       =			[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
            lowercase :       Optional[Any]       =			dummy_past_residuals[:]

            lowercase :       Optional[int]       =			scheduler.step_prk(_A				,						0				,						_A				,						**_A				).prev_sample
            lowercase :       str       =			scheduler.step_prk(_A				,						1				,						_A				,						**_A				).prev_sample

            self.assertEqual(output_a.shape				,						sample.shape				)
            self.assertEqual(output_a.shape				,						output_a.shape				)

            lowercase :       Dict       =			scheduler.step_plms(_A				,						0				,						_A				,						**_A				).prev_sample
            lowercase :       Tuple       =			scheduler.step_plms(_A				,						1				,						_A				,						**_A				).prev_sample

            self.assertEqual(output_a.shape				,						sample.shape				)
            self.assertEqual(output_a.shape				,						output_a.shape				)






  def  __a      (  self		:     Any				)		->   Optional[Any]:






       """simple docstring"""







       for timesteps in [100, 1_000]:
            self.check_over_configs(num_train_timesteps=_A				)






  def  __a      (  self		:     Tuple				)		->   Any:






       """simple docstring"""







       for steps_offset in [0, 1]:
            self.check_over_configs(steps_offset=_A				)

       lowercase :       str       =			self.scheduler_classes[0]
       lowercase :       int       =			self.get_scheduler_config(steps_offset=1				)
       lowercase :       Union[str, Any]       =			scheduler_class(**_A				)
       scheduler.set_timesteps(10				)
       assert torch.equal(
           scheduler.timesteps				,						torch.LongTensor(
               [901, 851, 851, 801, 801, 751, 751, 701, 701, 651, 651, 601, 601, 501, 401, 301, 201, 101, 1]				)				,						)






  def  __a      (  self		:     int				)		->   Optional[Any]:






       """simple docstring"""







       for beta_start, beta_end in zip([0.0_001, 0.001]				,						[0.002, 0.02]				):
            self.check_over_configs(beta_start=_A				,						beta_end=_A				)






  def  __a      (  self		:     List[str]				)		->   Tuple:






       """simple docstring"""







       for schedule in ["linear", "squaredcos_cap_v2"]:
            self.check_over_configs(beta_schedule=_A				)






  def  __a      (  self		:     List[str]				)		->   Optional[Any]:






       """simple docstring"""







       for prediction_type in ["epsilon", "v_prediction"]:
            self.check_over_configs(prediction_type=_A				)






  def  __a      (  self		:     Tuple				)		->   str:






       """simple docstring"""







       for t in [1, 5, 10]:
            self.check_over_forward(time_step=_A				)






  def  __a      (  self		:     List[Any]				)		->   str:






       """simple docstring"""







       for t, num_inference_steps in zip([1, 5, 10]				,						[10, 50, 100]				):
            self.check_over_forward(num_inference_steps=_A				)






  def  __a      (  self		:     str				)		->   List[Any]:






       """simple docstring"""







       lowercase :       str       =			27

       for scheduler_class in self.scheduler_classes:
            lowercase :       List[str]       =			self.dummy_sample
            lowercase :       Union[str, Any]       =			0.1 * sample

            lowercase :       int       =			self.get_scheduler_config()
            lowercase :       Union[str, Any]       =			scheduler_class(**_A				)

            scheduler.set_timesteps(_A				)

            # before power of 3 fix, would error on first step, so we only need to do two
            for i, t in enumerate(scheduler.prk_timesteps[:2]				):
                 lowercase :       Union[str, Any]       =			scheduler.step_prk(_A				,						_A				,						_A				).prev_sample






  def  __a      (  self		:     List[Any]				)		->   List[Any]:






       """simple docstring"""







       with self.assertRaises(_A				):
            lowercase :       Dict       =			self.scheduler_classes[0]
            lowercase :       Union[str, Any]       =			self.get_scheduler_config()
            lowercase :       Optional[int]       =			scheduler_class(**_A				)

            scheduler.step_plms(self.dummy_sample				,						1				,						self.dummy_sample				).prev_sample






  def  __a      (  self		:     str				)		->   Dict:






       """simple docstring"""







       lowercase :       List[str]       =			self.full_loop()
       lowercase :       List[Any]       =			torch.sum(torch.abs(_A				)				)
       lowercase :       Any       =			torch.mean(torch.abs(_A				)				)

       assert abs(result_sum.item() - 198.1_318				) < 1E-2
       assert abs(result_mean.item() - 0.2_580				) < 1E-3






  def  __a      (  self		:     Union[str, Any]				)		->   List[str]:






       """simple docstring"""







       lowercase :       List[str]       =			self.full_loop(prediction_type='''v_prediction'''				)
       lowercase :       Optional[int]       =			torch.sum(torch.abs(_A				)				)
       lowercase :       Tuple       =			torch.mean(torch.abs(_A				)				)

       assert abs(result_sum.item() - 67.3_986				) < 1E-2
       assert abs(result_mean.item() - 0.0_878				) < 1E-3






  def  __a      (  self		:     str				)		->   Optional[Any]:






       """simple docstring"""







       lowercase :       Any       =			self.full_loop(set_alpha_to_one=_A				,						beta_start=0.01				)
       lowercase :       List[str]       =			torch.sum(torch.abs(_A				)				)
       lowercase :       Optional[int]       =			torch.mean(torch.abs(_A				)				)

       assert abs(result_sum.item() - 230.0_399				) < 1E-2
       assert abs(result_mean.item() - 0.2_995				) < 1E-3






  def  __a      (  self		:     List[str]				)		->   Union[str, Any]:






       """simple docstring"""







       lowercase :       Union[str, Any]       =			self.full_loop(set_alpha_to_one=_A				,						beta_start=0.01				)
       lowercase :       int       =			torch.sum(torch.abs(_A				)				)
       lowercase :       int       =			torch.mean(torch.abs(_A				)				)

       assert abs(result_sum.item() - 186.9_482				) < 1E-2
       assert abs(result_mean.item() - 0.2_434				) < 1E-3